diff --git a/.gitignore b/.gitignore index b5f08c2..e6be4e7 100644 --- a/.gitignore +++ b/.gitignore @@ -6,3 +6,5 @@ dist/ build/ *.zip .pytest_cache/ + +sources/ diff --git a/README.md b/README.md index fc49c0e..6ac17e9 100644 --- a/README.md +++ b/README.md @@ -1,331 +1,64 @@ # SnapForm -`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) 背景剔除功能测试。 +SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一组序列帧快速变成有真实深度的立体面片。 -## 当前功能 +## 它能做什么 -- 使用 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，保留原图并换用新建的透明背景图片 -- 使用浮点深度贴图、Simple Subdivision、Displace Modifier 和原图 Bump - 构建立体面片 -- 转换对话框可选择推理模式、处理分辨率和面片网格细分 -- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹，不向当前 Blender 文件添加任何数据 -- 完整保留每次任务生成的所有官方文件 -- 安装和模型下载使用后台子进程；推理提交给常驻 AI Host，状态栏实时显示阶段、百分比进度条与取消按钮 -- 后台任务运行时，Install、Download 和 Export 按钮保持原位置并变灰，任务完成后恢复 +- 一键把 Image Empty 转换为立体面片：让 AI 估算深度，或直接用原图作为深度 +- 一键去除图片背景，生成透明前景图 +- 支持视频与序列帧输入，逐帧处理，动画也能得到逐帧对应的立体面片 +- 全部操作都在右键菜单里完成，不需要写任何代码 -## 使用 +## 安装 -1. 使用 `uv run python pack.py` 构建 Extension ZIP。 -2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。 -3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。 -4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn， - 并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进，每段进度 - 从 0 到 100%：下载 uv（如需）、创建 Python Runtime、下载 PyTorch、安装 xFormers（macOS 跳过）、 - 安装 AI 依赖、修补 DA3（跳过视频渲染），然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`， - 模型下载阶段会实时显示下载速度，例如 `6/8 da3 small: downloading 14.2MB/s`。 - Runtime 就绪后即使模型下载中断也不影响 Host 使用，可稍后在 `Models` 面板单独补下。 -5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 -6. 如需其他模型，展开 `Host > Models`，找到模型并点击该行的 `Download`。 -7. 展开默认收起的 `Debug`，选择 `Export Format` 和本次执行使用的模型。 -8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数，然后点击 `Export`。 -9. 测试背景剔除时，确认 `Host > Models` 中的 `BEN2 Base` 已就绪，然后在 `Debug` 的 - `BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。 - 完成后会打开 job 文件夹，其中的 `foreground.png` 带透明 Alpha；状态栏显示 - BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。 +1. 下载对应系统的扩展压缩包： +- Windows：`SnapForm.v0.10.10.windows-x64.zip` +- macOS（Apple 芯片）：`SnapForm.v0.10.10.macos-arm64.zip` +- macOS（Intel）：`SnapForm.v0.10.10.macos-x64.zip` +- Linux：`SnapForm.v0.10.10.linux-x64.zip` +2. 打开 Blender，进入 `Preferences > Extensions` +3. 点击 `Install from Disk`，选择下载的压缩包 +4. 在扩展偏好设置中确认 `uv` 路径与 Torch 后端（保持默认即可）；如果下载经常超时， + 可以在 `PyPI Index` 里填一个国内镜像，例如 `https://pypi.tuna.tsinghua.edu.cn/simple` +5. 打开 3D 视图右侧的 `SnapForm` 面板，点击 `Install Runtime` -## 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 Depth 的阶段包括 Runtime 启动、输入准备、模型加载、深度估算、缓存 -写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。安装中断后残留的 -半成品 Runtime 会在下次安装前自动清理，不会因为中断而无法再次安装。 +选中 Image Empty 后，右键选择 `Remove Background`，完成后 Empty 会自动换成透明前景图。 -导入过程精确保留 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` 仍用于 -确定原始相机面片，随后对象 XY Transform 将它匹配到源 Image Empty。生成对象会记录 -`snapform_reference_depth`、`snapform_camera_plane_width` 和 -`snapform_camera_plane_height`，便于检查实际采用的 -参考深度和尺寸。 +任务结果保存在 `~/.snapform/jobs/` 下，每个任务一个文件夹。视频或序列帧生成的面片会引用这些文件夹 +里的序列帧，请保留对应文件夹，不要随意删除。 -选择 `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 模型的深度不应 -被解释为现实世界的米；需要现实尺度时应选择带 metric depth 的 -`DA3 Nested Giant Large 1.1`。此外，当前结果仍是固定 XY 平面上的深度位移面，不是使用 -相机内参逐像素反投影得到的完整 camera-space 网格；中位深度截面为现有 Displace Modifier -工作流提供正确的整体高宽尺度。 +- 首次使用需要下载模型，如果下载慢或失败，稍后在 `Host > Models` 里重试即可 +- 安装 Runtime 时提示下载超时：新版会自动探测并切换到镜像、失败时自动重试一次；如果仍然 + 超时，在扩展偏好设置的 `PyPI Index` 中填入镜像地址（如清华）后重新安装即可 +- 找不到右键入口时，请确认选中了 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， -不会 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 -/ -└── scene.glb - -/ -└── gs_ply/*.ply - -/ -└── depth_vis/*.jpg - -/ -└── exports/mini_npz/results.npz - -/ -└── 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 下载前会在界面显示。 +面向开发者的技术内容见 [docs/technical.md](docs/technical.md)。 diff --git a/docs/technical.md b/docs/technical.md new file mode 100644 index 0000000..2378dc3 --- /dev/null +++ b/docs/technical.md @@ -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 下载前会在界面显示。 diff --git a/pack.py b/pack.py index f60ad12..ca4be0c 100644 --- a/pack.py +++ b/pack.py @@ -1,3 +1,5 @@ +import io +import tarfile import tomllib from pathlib import Path from zipfile import ZIP_DEFLATED, ZipFile @@ -5,30 +7,135 @@ from zipfile import ZIP_DEFLATED, ZipFile PROJECT_ROOT = Path(__file__).resolve().parent SOURCE_DIR = PROJECT_ROOT / "src" / "snapform" +SOURCES_DIR = PROJECT_ROOT / "sources" 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(): - manifest = tomllib.loads( - (SOURCE_DIR / "blender_manifest.toml").read_text(encoding="utf-8") +def bundled_source_bytes(filename): + """返回内置源码包内容；patches// 目录中的文件会覆盖官方包对应路径。""" + 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 / ( - f"{manifest['name']}.v{manifest['version']}.zip" + f"{manifest['name']}.v{manifest['version']}.{platform}.zip" ) 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: for source_path in sorted(SOURCE_DIR.rglob("*")): if not source_path.is_file(): continue if "__pycache__" in source_path.parts: continue + if source_path.name == "blender_manifest.toml": + continue 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 +def build_all(): + return [build_extension(platform) for platform in PLATFORMS] + + if __name__ == "__main__": - extension_path = build_extension() - size_mib = extension_path.stat().st_size / (1024 * 1024) - print(f"Built {extension_path} ({size_mib:.2f} MiB)") + extension_paths = build_all() + for extension_path in extension_paths: + size_mib = extension_path.stat().st_size / (1024 * 1024) + print(f"Built {extension_path} ({size_mib:.2f} MiB)") diff --git a/pyproject.toml b/pyproject.toml index 22114fe..8f2755f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "snapform" -version = "0.9.32" +version = "0.10.10" description = "Fast image-guided structure building tools for Blender" readme = "README.md" requires-python = ">=3.11" diff --git a/src/snapform/ai_host.py b/src/snapform/ai_host.py index be078d8..7f7a062 100644 --- a/src/snapform/ai_host.py +++ b/src/snapform/ai_host.py @@ -278,6 +278,11 @@ class HostState: model_cache=self.models, cancel_check=cancel_check, ) + elif payload["command"] == "extract-frames": + worker.extract_frames( + args, + cancel_check=cancel_check, + ) else: worker.reconstruct( args, @@ -426,7 +431,11 @@ def create_app(token, instance_id="test"): client_id = str(payload.get("client_id", "")) command = payload.get("command") 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") if not isinstance(parameters, dict) or not parameters.get("status"): raise HTTPException(status_code=422, detail="Missing job parameters") diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index d283249..11e1de7 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -1,7 +1,7 @@ schema_version = "1.0.0" id = "snapform" -version = "0.9.32" +version = "0.10.10" name = "SnapForm" tagline = "Build spatial structures from images, fast" maintainer = "OmooLab" diff --git a/src/snapform/bootstrap_runtime.py b/src/snapform/bootstrap_runtime.py index 68621ff..62d1416 100644 --- a/src/snapform/bootstrap_runtime.py +++ b/src/snapform/bootstrap_runtime.py @@ -20,21 +20,13 @@ if str(SCRIPT_DIRECTORY) not in sys.path: sys.path.insert(0, str(SCRIPT_DIRECTORY)) from models import BEN2_MODEL, MODELS -from worker import HF_ENDPOINT, HF_MIRRORS, format_rate, reachable -UV_VERSION = "0.11.3" -PREFLIGHT_MIN_FREE_BYTES = 8 * 1024**3 -PREFLIGHT_TIMEOUT = 8.0 -MIRROR_TIMEOUT = 5.0 PYPI_MIRRORS = ( "https://pypi.tuna.tsinghua.edu.cn/simple", "https://mirrors.aliyun.com/pypi/simple/", ) -TORCH_MIRRORS = { - "cu128": ("https://mirrors.aliyun.com/pytorch-wheels/cu128",), - "cpu": ("https://mirrors.aliyun.com/pytorch-wheels/cpu",), -} +PYPI_DEFAULT_INDEX = "https://pypi.org/simple/" TORCH_VERSION = "2.9.1" TORCHVISION_VERSION = "0.24.1" XFORMERS_VERSION = "0.0.33.post2" @@ -47,11 +39,7 @@ TORCH_INDEXES = { "cpu": "https://download.pytorch.org/whl/cpu", } _resolved_pypi_index = None -_resolved_hf_endpoint = None -_resolved_torch_index = None DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key) -PREFLIGHT_END = 0.05 -MIRROR_END = 0.07 RUNTIME_SLICE_START = 0.07 RUNTIME_SLICE_END = 0.88 VERIFY_SLICE_START = 0.88 @@ -101,8 +89,6 @@ def child_environment(environment=None): process_environment[env_key] = value if _resolved_pypi_index: process_environment["PIP_INDEX_URL"] = _resolved_pypi_index - if _resolved_hf_endpoint: - process_environment["HF_ENDPOINT"] = _resolved_hf_endpoint if environment: process_environment.update(environment) return process_environment @@ -131,127 +117,30 @@ def parse_package_count(line): def effective_torch_index(backend): - if _resolved_torch_index: - return _resolved_torch_index return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend) -def probe_hosts(urls, timeout): - results = {} - - def worker(url): - results[url] = reachable(url, timeout) - - threads = [ - threading.Thread(target=worker, args=(url,)) for url in urls - ] - for thread in threads: - thread.start() - for thread in threads: - thread.join() - return results - - -def resolve_endpoints(backend): - """Resolve package indexes, falling back to mirrors when direct access fails; returns (resolved, failures, used_mirror).""" - preferred = { - "pypi": os.environ.get("PIP_INDEX_URL") or "https://pypi.org/simple/", - "hf": os.environ.get("HF_ENDPOINT") or HF_ENDPOINT, - "torch": os.environ.get("SNAPFORM_TORCH_INDEX") - or TORCH_INDEXES.get(backend), - "github": "https://github.com/", - } - if sys.platform == "darwin": - preferred["torch"] = None - mirror_map = { - "pypi": PYPI_MIRRORS, - "hf": HF_MIRRORS, - "torch": TORCH_MIRRORS.get(backend, ()), - "github": (), - } - results = probe_hosts( - [url for url in preferred.values() if url], - PREFLIGHT_TIMEOUT, +def resolve_indexes(args): + """???????????????????????????????????""" + global _resolved_pypi_index + _resolved_pypi_index = ( + getattr(args, "pypi_index", None) + or os.environ.get("PIP_INDEX_URL") ) - resolved = {} - failures = [] - for key, url in preferred.items(): - if url is None: - resolved[key] = None - continue - if results[url]: - resolved[key] = url - continue - mirror_results = probe_hosts(mirror_map[key], MIRROR_TIMEOUT) - mirror = next( - ( - candidate - for candidate in mirror_map[key] - if mirror_results.get(candidate) - ), - None, - ) - if mirror: - resolved[key] = mirror - else: - resolved[key] = url - failures.append(url) - used_mirror = any( - key in preferred - and preferred[key] is not None - and resolved[key] != preferred[key] - for key in resolved - ) - return resolved, failures, used_mirror -def preflight(args): - """Check disk space and network reachability; switch to mirrors when direct access fails and raise a clear English error when nothing is reachable.""" - global _resolved_pypi_index, _resolved_hf_endpoint, _resolved_torch_index - - problems = [] - free_bytes = shutil.disk_usage(Path(args.runtime_dir).parent).free - if free_bytes < PREFLIGHT_MIN_FREE_BYTES: - needed_gb = PREFLIGHT_MIN_FREE_BYTES // (1024**3) - problems.append( - f"Not enough free disk space: {needed_gb} GB required, " - f"only {free_bytes / 1024**3:.1f} GB available." - ) - resolved, failures, used_mirror = resolve_endpoints(args.backend) - _resolved_pypi_index = resolved["pypi"] - _resolved_hf_endpoint = resolved["hf"] - _resolved_torch_index = resolved["torch"] - unreachable = failures - if unreachable: - proxies = urllib.request.getproxies() - proxy_hint = ( - "Proxy detected: " + ", ".join( - f"{key}={value}" for key, value in proxies.items() - ) - if proxies - else "No proxy detected" - ) - problems.append( - "Cannot reach: " + ", ".join(unreachable) - + f". {proxy_hint}. Check network or proxy settings and retry." - ) - if problems: - raise RuntimeError("Prerequisites check failed: " + " ".join(problems)) - return resolved, used_mirror - - -def run_install( +def _run_install_once( status_path, stage, stages, label, command, - environment=None, - start=0.0, - end=1.0, - clock=time.monotonic, + environment, + start, + end, + 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])) process_environment = child_environment(environment) process_environment["UV_LOG"] = "debug" @@ -342,6 +231,42 @@ def run_install( 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): if sys.platform == "win32": return Path(runtime_dir) / "Scripts" / "python.exe" @@ -355,15 +280,22 @@ def install_stages(args, runtime_dir, python, uv): if sys.platform == "darwin" else effective_torch_index(args.backend) ) - def github_archive_url(repository, commit): - return f"https://github.com/{repository}/archive/{commit}.tar.gz" + def archive_requirement(project_name, filename): + 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 = ( - "depth-anything-3 @ " - + github_archive_url("ByteDance-Seed/Depth-Anything-3", args.da3_commit) + da3_requirement = archive_requirement( + "depth-anything-3", + "depth-anything-3.tar.gz", ) - ben2_requirement = ( - "ben2 @ " + github_archive_url("PramaLLC/BEN2", args.ben2_commit) + ben2_requirement = archive_requirement( + "ben2", + "ben2.tar.gz", ) torch_command = [ uv, @@ -488,12 +420,41 @@ def uv_download_target(): 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() - return ( - "https://github.com/astral-sh/uv/releases/download/" - f"{UV_VERSION}/uv-{target}.{archive}" - ) + path = bundled_sources_directory() / f"uv-{target}.{archive}" + return path if path.is_file() else None + + +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): @@ -511,8 +472,8 @@ def extract_uv(archive_path, uv_path): os.chmod(uv_path, 0o755) -def download_uv( - uv_path, +def install_uv( + uv, status_path, stage=None, stages=None, @@ -520,50 +481,29 @@ def download_uv( start=0.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": stage, "stages": stages, "stage_label": stage_label, } - - def write(progress, message): - write_status( - status_path, - start + progress * (end - start), - message, - **stage_fields, + bundled_uv = bundled_uv_archive() + if bundled_uv is None: + raise RuntimeError( + "Bundled uv is missing; reinstall the extension matching your " + f"platform (searched {bundled_sources_directory()})" ) - - write(0.0, "downloading") - uv_path = Path(uv_path) - uv_path.parent.mkdir(parents=True, exist_ok=True) - archive_path = uv_path.with_name(f"{uv_path.name}.part") - request = urllib.request.Request( - uv_download_url(), - headers={"User-Agent": "SnapForm-Runtime-Installer"}, + if not archive_is_valid(bundled_uv): + raise RuntimeError(f"Bundled uv archive is invalid: {bundled_uv}") + print(f"Extracting bundled uv from {bundled_uv}") + write_status( + status_path, + start, + "extracting bundled uv", + **stage_fields, ) - try: - with urllib.request.urlopen(request, timeout=60) as response: - total = int(response.headers.get("Content-Length") or 0) - downloaded = 0 - started = time.monotonic() - with archive_path.open("wb") as output: - while True: - chunk = response.read(65536) - if not chunk: - break - output.write(chunk) - downloaded += len(chunk) - progress = ( - min(downloaded / total, 1.0) if total else 0.0 - ) - speed = downloaded / max(time.monotonic() - started, 1e-6) - write(progress, f"downloading {format_rate(speed)}") - write(1.0, "preparing") - extract_uv(archive_path, uv_path) - finally: - archive_path.unlink(missing_ok=True) + Path(uv).parent.mkdir(parents=True, exist_ok=True) + extract_uv(bundled_uv, uv) def run_stage( @@ -636,10 +576,7 @@ def install(args): "SnapForm requires Apple Silicon macOS; PyTorch stopped " "shipping Intel Mac (x86_64) wheels after 2.2.2" ) - write_status(args.status, 0.0, "checking prerequisites") - _resolved, used_mirror = preflight(args) - if used_mirror: - write_status(args.status, MIRROR_END, "using mirror sources") + resolve_indexes(args) runtime_dir = Path(args.runtime_dir) runtime_dir.parent.mkdir(parents=True, exist_ok=True) if not runtime_ready(runtime_dir): @@ -650,6 +587,10 @@ def install(args): uv = Path(resolved) if resolved else uv uv_missing = not uv.is_file() 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) runtime_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 ) if uv_missing and index == 1: - download_uv( + install_uv( uv, args.status, stage=index, @@ -766,6 +707,7 @@ def parse_args(): parser.add_argument("--da3-commit", required=True) parser.add_argument("--ben2-commit", required=True) parser.add_argument("--status", required=True) + parser.add_argument("--pypi-index", default=None) return parser.parse_args() diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index d71ec2e..6228d53 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -1,3 +1,5 @@ +import re +import shutil from pathlib import Path import bpy @@ -22,6 +24,17 @@ IMAGE_NAME_SUFFIXES = ( ".bmp", ".tga", ) +VIDEO_NAME_SUFFIXES = ( + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", +) 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): filepath_from_user = getattr(image, "filepath_from_user", None) filepath = ( @@ -65,10 +128,22 @@ def image_is_packed(image): 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): name = image.name 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) if suffix_index <= 0: continue @@ -85,8 +160,15 @@ def depth_data_name(source_object): 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): - return f"{image_base_name(source_object.data)}.rb" + return f"{image_base_name(source_object.data)}.foreground" def scaled_process_resolution(image, scale): @@ -117,10 +199,60 @@ def save_packed_image(image, scene): 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.*?)(?P\d+)(?P\.[^.]+)$", + 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.*?)(?P\d+)(?P\.[^.]+)$", + 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): source_path = image_path(image) if image_is_packed(image): 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(): raise RuntimeError( "The image empty must reference an existing image file" @@ -128,12 +260,207 @@ def prepare_image_input(image, scene): 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( source_object, output_path, assign_to_source=True, ): 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(): raise RuntimeError( 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)) -def extract_depth_array(npz_path): +def extract_depth_arrays(npz_path): try: import numpy as np except ImportError as error: @@ -256,13 +583,17 @@ def extract_depth_array(npz_path): if depth.ndim >= 3 and depth.shape[-1] == 1: depth = depth[..., 0] - if depth.ndim == 3: - depth = depth[0] - if depth.ndim != 2: + if depth.ndim == 2: + depth = depth[None, ...] + if depth.ndim != 3: raise ValueError( 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): @@ -475,6 +806,13 @@ def create_material( displacement_midlevel=0.0, displacement_scale=1.0, 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.use_nodes = True @@ -491,6 +829,25 @@ def create_material( color_texture.extension = "EXTEND" color_texture.location = (-640.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 if unshaded: @@ -558,6 +915,14 @@ def create_material( depth_texture.image = displacement_image depth_texture.interpolation = "Linear" 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_midlevel ) @@ -612,8 +977,35 @@ def build_image_depth_plane( displacement_via="MODIFIER", unshaded=False, color_image=None, + depth_frames_directory=None, ): 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( source_object, subdivision_level=subdivision_level, @@ -632,25 +1024,43 @@ def build_image_depth_plane( depth_data_name(source_object), 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.texture = texture displace.texture_coords = "UV" displace.mid_level = ORIGINAL_DISPLACEMENT_MIDLEVEL depth_scale = abs(float(plane_object.scale.z)) + displacement_image = ( + depth_source_image if displacement_via == "MATERIAL" else None + ) material = create_material( source_name, color_image if color_image is not None else source_object.data, depth_scale, - displacement_image=( - source_object.data if displacement_via == "MATERIAL" else None - ), + displacement_image=displacement_image, displacement_midlevel=ORIGINAL_DISPLACEMENT_MIDLEVEL, displacement_scale=ORIGINAL_DISPLACEMENT_STRENGTH, 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) + plane_object["snapform_frame_count"] = frame_count return finish_depth_plane( context, source_object, @@ -669,9 +1079,11 @@ def build_depth_plane( unshaded=False, depth_name="", color_image=None, + depth_frames_directory=None, ): 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) depth_values = prepare_depth_values(depth) bounds, reference_depth = camera_plane_bounds( @@ -681,10 +1093,25 @@ def build_depth_plane( DISPLACEMENT_STRENGTH, ) 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_name, 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) plane_object = bpy.data.objects.new( source_name, @@ -709,6 +1136,13 @@ def build_depth_plane( ) texture.image = depth_image 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.texture = texture displace.texture_coords = "UV" @@ -726,12 +1160,22 @@ def build_depth_plane( displacement_midlevel=DISPLACEMENT_MIDLEVEL, displacement_scale=DISPLACEMENT_STRENGTH, 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) plane_object["snapform_reference_depth"] = reference_depth plane_object["snapform_camera_plane_width"] = bounds[1] - bounds[0] plane_object["snapform_camera_plane_height"] = bounds[3] - bounds[2] plane_object["snapform_depth_scale"] = depth_scale + plane_object["snapform_frame_count"] = frame_count return finish_depth_plane( context, @@ -758,6 +1202,12 @@ class BuildDepthPlane(bpy.types.Operator): snapform_background_path: bpy.props.StringProperty( 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( default=5, min=0, @@ -806,6 +1256,11 @@ class BuildDepthPlane(bpy.types.Operator): 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_depth_plane( context, source_object, @@ -816,6 +1271,9 @@ class BuildDepthPlane(bpy.types.Operator): unshaded=self.snapform_unshaded, depth_name=self.snapform_depth_data_name, color_image=color_image, + depth_frames_directory=( + self.snapform_depth_frames_dir or None + ), ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -834,6 +1292,12 @@ class BuildOriginalDepthPlane(bpy.types.Operator): snapform_background_path: bpy.props.StringProperty( 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( default=5, min=0, @@ -866,11 +1330,18 @@ class BuildOriginalDepthPlane(bpy.types.Operator): ) return {"CANCELLED"} try: - color_image = create_removed_background_image( - source_object, - self.snapform_background_path, - assign_to_source=False, - ) + color_image = None + if self.snapform_background_path: + color_image = create_removed_background_image( + 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( context, source_object, @@ -879,6 +1350,9 @@ class BuildOriginalDepthPlane(bpy.types.Operator): displacement_via=self.snapform_displacement_via, unshaded=self.snapform_unshaded, color_image=color_image, + depth_frames_directory=( + self.snapform_depth_frames_dir or None + ), ) except (OSError, RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) @@ -938,6 +1412,11 @@ class RemoveImageBackground(bpy.types.Operator): from .operators import IMAGE_BACKGROUND_JOB 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: source_path, delete_input = prepare_image_input( source_object.data, @@ -956,7 +1435,7 @@ class RemoveImageBackground(bpy.types.Operator): if "CANCELLED" not in result: return {"FINISHED"} if delete_input: - source_path.unlink(missing_ok=True) + remove_temporary_input(source_path) return {"CANCELLED"} @@ -1089,13 +1568,45 @@ class ConvertImageToDepthPlane(bpy.types.Operator): layout.prop(self, "snapform_delete_original") 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 + if video_total_frames(source_object.data) > 1: + context.scene.snapform_settings.video_frames = ( + default_extract_frames(source_object) + ) if ( self.snapform_depth_source == "ORIGINAL" 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: build_image_depth_plane( context, @@ -1141,7 +1652,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): if "CANCELLED" not in result: return {"FINISHED"} if delete_input: - source_path.unlink(missing_ok=True) + remove_temporary_input(source_path) return {"CANCELLED"} image = source_object.data @@ -1179,7 +1690,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): ) if "CANCELLED" in result: if delete_input: - source_path.unlink(missing_ok=True) + remove_temporary_input(source_path) return {"CANCELLED"} global _last_depth_plane_model _last_depth_plane_model = self.snapform_model diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 2c4f09e..6599dad 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -1,3 +1,4 @@ +import shutil import subprocess from dataclasses import dataclass from pathlib import Path @@ -18,8 +19,35 @@ IMAGE_DEPTH_JOB = "IMAGE_DEPTH" BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND" +ORIGINAL_FRAMES_JOB = "ORIGINAL_FRAMES" IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 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): @@ -81,8 +109,12 @@ def close_active_job(): if _active_job.log_file is not None: _active_job.log_file.close() if _active_job.temporary_input_path is not None: + temporary_path = _active_job.temporary_input_path 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: pass _active_job = None @@ -102,6 +134,7 @@ class RunJob(bpy.types.Operator): (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), (ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""), + (ORIGINAL_FRAMES_JOB, "Extract Original Depth Frames", ""), ), options={"HIDDEN", "SKIP_SAVE"}, ) @@ -227,6 +260,7 @@ class RunJob(bpy.types.Operator): IMAGE_DEPTH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, + ORIGINAL_FRAMES_JOB, } and self.snapform_delete_input else None @@ -267,6 +301,7 @@ class RunJob(bpy.types.Operator): IMAGE_DEPTH_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, + ORIGINAL_FRAMES_JOB, }: configured_input = self.snapform_input_path elif self.job_type == BACKGROUND_REMOVAL_JOB: @@ -274,20 +309,24 @@ class RunJob(bpy.types.Operator): else: configured_input = settings.input_path input_path = Path(bpy.path.abspath(configured_input)) - if not input_path.is_file(): - raise ValueError("Select a valid image or video file") + if not (input_path.is_file() or input_path.is_dir()): + 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 { BACKGROUND_REMOVAL_JOB, IMAGE_BACKGROUND_JOB, ORIGINAL_BACKGROUND_JOB, }: - if input_path.suffix.lower() not in { - ".jpg", - ".jpeg", - ".png", - ".webp", - }: - raise ValueError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") + if not is_ben2_input(input_path): + raise ValueError( + "BEN2 requires an image, video, or directory of image frames" + ) if not runtime.model_ready(BEN2_MODEL.key): raise RuntimeError(f"Download {BEN2_MODEL.label} first") return ( @@ -297,6 +336,18 @@ class RunJob(bpy.types.Operator): paths["status"], paths["result"], 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, ) @@ -342,6 +393,7 @@ class RunJob(bpy.types.Operator): else None ), use_ray_pose=use_ray_pose, + start_frame=start_frame, remove_background=( self.job_type == IMAGE_DEPTH_JOB and remove_background @@ -371,6 +423,8 @@ class RunJob(bpy.types.Operator): return "Starting Image Empty background removal" if self.job_type == ORIGINAL_BACKGROUND_JOB: return "Starting background removal" + if self.job_type == ORIGINAL_FRAMES_JOB: + return "Extracting source frames" if self.job_type == "DOWNLOAD_MODEL": model = get_downloadable_model(self.model_key) return f"Starting {model.label} download" @@ -484,7 +538,7 @@ class RunJob(bpy.types.Operator): self._redraw_ui(context, force=True) return {"CANCELLED"} status["message"] = "Depth plane is ready" - if job.kind == ORIGINAL_BACKGROUND_JOB: + if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}: try: settings.status = "Building depth plane" settings.progress = 0.96 @@ -515,24 +569,39 @@ class RunJob(bpy.types.Operator): settings.status = status.get("message", "Task complete") settings.progress = 1.0 + try: + job.status_path.unlink(missing_ok=True) + except OSError: + pass close_active_job() self._redraw_ui(context, force=True) return {"FINISHED"} def _build_depth_plane(self, context, job): - npz_path = ( - job.result_directory - / "exports" - / "mini_npz" - / "results.npz" - ) + result_directory = job.result_directory + npz_path = result_directory / "results.npz" + is_sequence = bool(list(result_directory.glob("frame.*"))) + foreground_paths = sorted(result_directory.glob("foreground.*.png")) + 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( snapform_source_object_name=job.source_object_name, snapform_npz_path=str(npz_path), snapform_depth_data_name=job.depth_data_name, - snapform_background_path=( - str(job.result_directory / "foreground.png") - if job.remove_background + snapform_background_path=background_path, + snapform_color_frames_dir=( + str(result_directory) + if is_sequence and not foreground_paths + else "" + ), + snapform_depth_frames_dir=( + str(result_directory) + if is_sequence else "" ), 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") 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( snapform_source_object_name=job.source_object_name, - snapform_background_path=str( - job.result_directory / "foreground.png" + snapform_background_path=background_path, + 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_delete_source=job.delete_source, @@ -564,7 +647,12 @@ class RunJob(bpy.types.Operator): raise RuntimeError("Unable to create the depth plane Undo step") 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( snapform_source_object_name=job.source_object_name, snapform_output_path=str(output_path), diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 74e5d92..3cc7d66 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -35,6 +35,14 @@ class SnapFormPreferences(bpy.types.AddonPreferences): ), 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): from .operators import active_job @@ -44,6 +52,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences): layout.prop(self, "storage_root") layout.prop(self, "uv_executable") layout.prop(self, "torch_backend") + layout.prop(self, "pypi_index_url") layout.separator() layout.label(text=f"Runtime: {runtime_directory()}") if runtime_ready(): diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 9e5d478..6ade35b 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -15,8 +15,8 @@ class SnapFormSettings(bpy.types.PropertyGroup): subtype="FILE_PATH", ) video_frames: bpy.props.IntProperty( - name="Video Frames", - description="Maximum number of frames sampled evenly from a video", + name="Extract Frames", + description="Maximum number of consecutive frames extracted from the start of a video or frame sequence", default=16, min=2, soft_max=64, diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index cb0e93a..561792e 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -24,6 +24,13 @@ RUNTIME_SCHEMA = 4 HOST_SCHEMA = 2 HOST_ADDRESS = "127.0.0.1" 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) @@ -122,15 +129,26 @@ def temporary_input_path(): 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): timestamp = datetime.now().strftime("%Y%m%d-%H%M%S") - token = uuid.uuid4().hex[:8] - stem = f"{timestamp}-{job_type.lower()}-{token}" - directory = jobs_directory() + stem = f"{job_folder_name(job_type)}_{timestamp}" + directory = jobs_directory() / stem + directory.mkdir(parents=True, exist_ok=True) return { - "log": directory / f"{stem}.log", - "status": directory / f"{stem}.json", - "result": directory / stem, + "log": directory / f"{directory.name}.log", + "status": directory / f"{directory.name}.json", + "result": directory, } @@ -155,8 +173,9 @@ def install_runtime_command(status_path): executable = uv_executable() addon_preferences = preferences() backend = getattr(addon_preferences, "torch_backend", "CU128") + pypi_index = getattr(addon_preferences, "pypi_index_url", "") or "" script = Path(__file__).with_name("bootstrap_runtime.py") - return [ + command = [ sys.executable, str(script), "--uv", @@ -174,6 +193,9 @@ def install_runtime_command(status_path): "--status", str(status_path), ] + if pypi_index: + command.extend(["--pypi-index", pypi_index]) + return command def host_log_path(): @@ -616,6 +638,8 @@ def remove_background_request( status_path, result_directory, model_key, + video_frames=None, + start_frame=1, ): return HostJobRequest( { @@ -627,6 +651,39 @@ def remove_background_request( "output_dir": str(result_directory), "status": str(status_path), "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, use_ray_pose=None, remove_background=False, + start_frame=1, ): snapform_selected_key = model_key or selected_model_key(settings, export_format) snapform_selected_process_res = ( @@ -673,6 +731,7 @@ def reconstruct_request( "show_cameras": 1 if settings.show_cameras else 0, "gs_views_interval": settings.gs_views_interval, "video_frames": settings.video_frames, + "start_frame": start_frame, "device": settings.device.lower(), "use_ray_pose": 1 if snapform_selected_ray_pose else 0, "ref_view_strategy": settings.ref_view_strategy, diff --git a/src/snapform/worker.py b/src/snapform/worker.py index bc09b79..f2fc3b6 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -1,13 +1,14 @@ import argparse +import fnmatch import json import os import shutil import tempfile import threading import time -import urllib.error -import urllib.request from pathlib import Path +from urllib.parse import quote, urlencode +from urllib.request import Request, urlopen IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} @@ -24,9 +25,10 @@ VIDEO_SUFFIXES = { } SCENE_FILENAME = "scene.glb" EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz") -HF_ENDPOINT = "https://huggingface.co/" -HF_MIRRORS = ("https://hf-mirror.com",) -HF_PROBE_TIMEOUT = 5.0 +HF_ENDPOINT = "https://huggingface.co" +MODELSCOPE_ENDPOINT = "https://modelscope.cn" +MODELSCOPE_REVISION = "master" +MODELSCOPE_TIMEOUT = 60.0 class TaskCancelled(RuntimeError): @@ -38,34 +40,6 @@ def check_cancelled(cancel_check): raise TaskCancelled("Task cancelled") -def reachable(url, timeout=HF_PROBE_TIMEOUT): - request = urllib.request.Request( - url, - method="HEAD", - headers={"User-Agent": "SnapForm-Runtime-Installer"}, - ) - try: - with urllib.request.urlopen(request, timeout=timeout): - return True - except urllib.error.HTTPError: - return True - except Exception: - return False - - -def resolve_hf_endpoint(): - """Probe Hugging Face reachability and fall back to a mirror when direct access fails; returns (endpoint, used_mirror).""" - configured = os.environ.get("HF_ENDPOINT") - if configured: - return configured, False - if reachable(HF_ENDPOINT): - return HF_ENDPOINT, False - for mirror in HF_MIRRORS: - if reachable(mirror): - return mirror, True - return HF_ENDPOINT, False - - def write_status( path, progress, @@ -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( video_path, output_directory, maximum_frames, cv2, + start_frame=1, cancel_check=None, ): + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) capture = cv2.VideoCapture(str(video_path)) if not capture.isOpened(): 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 = [] + skipped = 0 try: - for output_index, frame_index in enumerate(frame_indices): + while len(image_paths) < maximum_frames: check_cancelled(cancel_check) - capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index)) success, frame = capture.read() if not success: - raise RuntimeError( - f"Unable to read video frame {int(frame_index)}" - ) - image_path = output_directory / f"frame-{output_index:04d}.jpg" + break + if skipped < start_frame - 1: + skipped += 1 + continue + image_path = ( + output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" + ) if not cv2.imwrite(str(image_path), frame): raise RuntimeError(f"Unable to write video frame: {image_path}") image_paths.append(image_path) finally: capture.release() + if not image_paths: + raise RuntimeError(f"Unable to read any video frame from: {video_path}") 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) 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() if suffix in IMAGE_SUFFIXES: return [input_path], None @@ -174,6 +158,7 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None): temporary_directory, video_frames, cv2, + start_frame=start_frame, cancel_check=cancel_check, ) except Exception: @@ -182,6 +167,43 @@ def prepare_input(input_path, video_frames, temporary_root, cancel_check=None): 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: def __init__( self, @@ -318,21 +340,8 @@ def download_model(args): "stages": getattr(args, "stages", 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 - message = f"downloading {label}" - if used_mirror: - message = f"{message} (mirror)" - write_status(args.status, start, message, **stage_kwargs) progress_class = create_download_progress_class( args.status, args.repository, @@ -341,18 +350,213 @@ def download_model(args): label=label, **stage_kwargs, ) - snapshot_download( - repo_id=args.repository, - local_dir=model_dir, - tqdm_class=progress_class, - allow_patterns=args.allow_pattern or None, + failures = [] + for endpoint in download_endpoints(): + if not hf_endpoint_reachable(endpoint): + failures.append((endpoint, "endpoint unreachable")) + 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, - start + span, - "model ready", - **stage_kwargs, + repository, + start=start, + 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): @@ -435,6 +639,49 @@ def infer_foreground( 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): total_started = time.perf_counter() 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 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) - write_status(args.status, 0.12, f"Loading BEN2 on {device.type}") - write_status(args.status, 0.35, "Removing background") - foreground, load_ms, inference_ms = infer_foreground( + video_frames = getattr(args, "video_frames", 16) + start_frame = getattr(args, "start_frame", 1) + images, temporary_directory = prepare_background_input( args.input, - args.model_dir, - device, - torch, - model_cache=model_cache, + video_frames, + output_directory.parent, + start_frame=start_frame, 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") + try: + write_status(args.status, 0.12, f"Loading BEN2 on {device.type}") + if len(images) == 1: + write_status(args.status, 0.35, "Removing background") + foreground, load_ms, inference_ms = infer_foreground( + images[0], + args.model_dir, + device, + torch, + model_cache=model_cache, + cancel_check=cancel_check, + ) + 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) total_ms = (time.perf_counter() - total_started) * 1000.0 + frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)" 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)" ) print(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): output_directory = Path(output_directory) if export_format == "glb": @@ -491,7 +787,12 @@ def exported_paths(output_directory, export_format): if export_format == "depth_vis": return sorted((output_directory / "depth_vis").glob("*.jpg")) 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 [] @@ -543,12 +844,13 @@ def reconstruct(args, model_cache=None, cancel_check=None): from depth_anything_3.utils.export import export 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") images, temporary_directory = prepare_input( args.input, args.video_frames, output_directory.parent, + start_frame=getattr(args, "start_frame", 1), cancel_check=cancel_check, ) @@ -562,22 +864,36 @@ def reconstruct(args, model_cache=None, cancel_check=None): ) if remove_background_model_dir: write_status(args.status, 0.12, "Removing image background") - foreground, load_ms, inference_ms = infer_foreground( - images[0], - remove_background_model_dir, - device, - torch, - model_cache=model_cache, - cancel_check=cancel_check, - ) - check_cancelled(cancel_check) - foreground.convert("RGBA").save( - output_directory / "foreground.png", - format="PNG", - compress_level=1, - ) - if not (output_directory / "foreground.png").is_file(): - raise RuntimeError("BEN2 did not write a foreground image") + if len(images) == 1: + foreground, load_ms, inference_ms = infer_foreground( + images[0], + remove_background_model_dir, + device, + torch, + model_cache=model_cache, + cancel_check=cancel_check, + ) + check_cancelled(cancel_check) + foreground.convert("RGBA").save( + output_directory / "foreground.png", + format="PNG", + compress_level=1, + ) + 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( f"BEN2 inference {inference_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( 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) + copy_frames(images, output_directory) write_status( args.status, 1.0, @@ -673,6 +1003,15 @@ def parser(): choices=("auto", "cuda", "mps", "cpu"), 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.add_argument("--input", required=True) @@ -709,6 +1048,7 @@ def parser(): required=True, ) reconstruction.add_argument("--video-frames", type=int, required=True) + reconstruction.add_argument("--start-frame", type=int, default=1) reconstruction.add_argument( "--device", choices=("auto", "cuda", "mps", "cpu"), @@ -735,6 +1075,8 @@ def main(): download_model(args) elif args.command == "reconstruct": reconstruct(args) + elif args.command == "extract-frames": + extract_frames(args) else: remove_background(args) except Exception as error: diff --git a/tests/test_addon.py b/tests/test_addon.py index c70e2a0..9f47880 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -1,5 +1,4 @@ import importlib -import inspect import sys import unittest from pathlib import Path @@ -52,6 +51,9 @@ class AddonTest(unittest.TestCase): ) fake_bpy.app = SimpleNamespace(online_access=True) fake_bpy.path = SimpleNamespace(abspath=lambda path: path) + fake_bpy.data = SimpleNamespace( + objects=SimpleNamespace(get=lambda _name: None), + ) sys.modules["bpy"] = fake_bpy sys.path.insert(0, str(SOURCE_ROOT)) cls.addon = importlib.import_module("snapform") @@ -75,6 +77,9 @@ class AddonTest(unittest.TestCase): self.context_menu_draws.clear() self.addon.operators._active_job = None self.addon.image_depth._last_depth_plane_model = "BASE" + self.fake_bpy.data = SimpleNamespace( + objects=SimpleNamespace(get=lambda _name: None), + ) if hasattr(self.fake_bpy.types.Scene, "snapform_settings"): del self.fake_bpy.types.Scene.snapform_settings @@ -127,43 +132,6 @@ class AddonTest(unittest.TestCase): {"INTERNAL"}, ) - def test_completed_depth_job_uses_undoable_build_operator(self): - source = inspect.getsource( - self.addon.RunJob._build_depth_plane - ) - - self.assertIn("bpy.ops.snapform.build_depth_plane(", source) - self.assertIn("snapform_unshaded=job.unshaded", source) - self.assertIn( - "snapform_depth_data_name=job.depth_data_name", - source, - ) - self.assertIn('if "FINISHED" not in result:', source) - self.assertIn( - 'bpy.ops.ed.undo_push(message="Convert to Depth Plane")', - source, - ) - self.assertIn('if "FINISHED" not in undo_result:', source) - - def test_image_context_menu_exposes_depth_and_background_actions(self): - draw_source = inspect.getsource( - self.addon.draw_image_context_menu - ) - - self.assertIn( - '"snapform.convert_image_to_depth_plane"', - draw_source, - ) - self.assertIn( - '"snapform.remove_image_background"', - draw_source, - ) - self.assertIn('text="Remove Background"', draw_source) - self.assertEqual(draw_source.count("row = self.layout.row()"), 2) - self.assertIn('operator_context = "INVOKE_DEFAULT"', draw_source) - self.assertNotIn("import_depth_plane", draw_source) - self.assertNotIn("Rebuild", draw_source) - def test_background_result_creates_new_packed_image(self): original_image = SimpleNamespace( name="img.png", @@ -211,7 +179,7 @@ class AddonTest(unittest.TestCase): source_object.empty_image_offset, (-0.25, -0.75), ) - self.assertEqual(result_image.name, "img.rb") + self.assertEqual(result_image.name, "img.foreground") self.assertEqual(pack_calls, [True]) self.assertEqual(remove_calls, []) self.assertEqual( @@ -228,7 +196,7 @@ class AddonTest(unittest.TestCase): self.addon.image_depth.removed_background_data_name( source_object ), - "img.001.rb", + "img.001.foreground", ) def test_depth_background_image_does_not_replace_source_empty(self): @@ -261,15 +229,67 @@ class AddonTest(unittest.TestCase): self.assertIs(created, result_image) self.assertIs(source_object.data, original_image) self.assertFalse(original_image.use_fake_user) - self.assertEqual(result_image.name, "img.rb") + self.assertEqual(result_image.name, "img.foreground") self.assertEqual(pack_calls, [True]) + def test_sequence_background_result_preserves_empty_image_start(self): + original_image = SimpleNamespace( + name="img.png", + use_fake_user=False, + ) + image_user = SimpleNamespace( + frame_start=4, + frame_offset=0, + frame_duration=1, + use_auto_refresh=False, + use_cyclic=False, + ) + source_object = SimpleNamespace( + data=original_image, + image_user=image_user, + ) + result_image = SimpleNamespace( + name="foreground.png", + source="SEQUENCE", + pack=lambda: None, + unpack=lambda **_options: None, + packed_files=(), + ) + self.fake_bpy.data = SimpleNamespace( + images=SimpleNamespace( + load=lambda *_args, **_options: result_image, + remove=lambda _image: None, + ), + ) + + with TemporaryDirectory() as directory: + result_directory = Path(directory) + (result_directory / "foreground.0001.png").touch() + (result_directory / "foreground.0002.png").touch() + created = self.addon.image_depth.create_removed_background_image( + source_object, + result_directory, + ) + + self.assertIs(created, result_image) + self.assertIs(source_object.data, result_image) + self.assertEqual(image_user.frame_start, 4) + self.assertEqual(image_user.frame_duration, 2) + self.assertEqual(image_user.frame_offset, 0) + self.assertTrue(image_user.use_auto_refresh) + self.assertFalse(image_user.use_cyclic) + def test_image_empty_background_action_starts_production_job(self): source_object = SimpleNamespace( name="Reference", data=SimpleNamespace(name="img.png"), ) - context = SimpleNamespace(object=source_object, scene=object()) + context = SimpleNamespace( + object=source_object, + scene=SimpleNamespace( + snapform_settings=SimpleNamespace(video_frames=16) + ), + ) calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( @@ -333,6 +353,52 @@ class AddonTest(unittest.TestCase): self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND") self.assertFalse(calls[0]["snapform_delete_source"]) + def test_original_depth_animation_extracts_frames_without_background(self): + operator = self.addon.ConvertImageToDepthPlane() + operator.snapform_depth_source = "ORIGINAL" + operator.snapform_remove_background = False + operator.snapform_displacement_via = "MODIFIER" + operator.snapform_unshaded = False + operator.snapform_subdivision_level = 5 + operator.snapform_delete_original = False + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace( + name="clip.mp4", + source="MOVIE", + frame_duration=97, + ), + ) + scene = SimpleNamespace( + snapform_settings=SimpleNamespace(video_frames=16) + ) + context = SimpleNamespace(object=source_object, scene=scene) + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "clip.mp4" + input_path.touch() + with patch.object( + self.addon.image_depth, + "prepare_image_input", + return_value=(input_path, False), + ): + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + self.assertEqual(scene.snapform_settings.video_frames, 97) + self.assertEqual(calls[0]["job_type"], "ORIGINAL_FRAMES") + self.assertEqual(calls[0]["snapform_input_path"], str(input_path)) + self.assertEqual(calls[0]["snapform_source_object_name"], "Reference") + self.assertFalse(calls[0]["snapform_delete_input"]) + self.assertFalse(calls[0]["snapform_delete_source"]) + def test_conversion_dialog_exposes_requested_settings(self): operator = self.addon.ConvertImageToDepthPlane @@ -393,50 +459,6 @@ class AddonTest(unittest.TestCase): operator.__annotations__["snapform_unshaded"]["name"], "Unshaded", ) - draw_source = inspect.getsource(operator.draw) - self.assertIn( - 'if self.snapform_depth_source == "ESTIMATE":', - draw_source, - ) - self.assertLess( - draw_source.index( - '"snapform_process_res_factor",' - ), - draw_source.index( - 'settings_body.prop(self, "snapform_subdivision_level")' - ), - ) - self.assertIn( - 'text=f"Resolution Scale ({process_res} px)"', - draw_source, - ) - self.assertIn("estimate_box = layout.box()", draw_source) - self.assertIn( - 'estimate_box.prop(self, "snapform_model")', - draw_source, - ) - self.assertIn( - 'estimate_box.prop(self, "snapform_use_ray_pose")', - draw_source, - ) - self.assertIn( - 'default_closed=True', - draw_source, - ) - self.assertIn('settings_header.label(text="Options")', draw_source) - self.assertIn('settings_header.label(text="Options")', draw_source) - self.assertNotIn("layout.separator", draw_source) - setting_fields = ( - "snapform_displacement_via", - "snapform_subdivision_level", - "snapform_remove_background", - "snapform_unshaded", - ) - field_positions = [ - draw_source.index(f'settings_body.prop(self, "{field}")') - for field in setting_fields - ] - self.assertEqual(field_positions, sorted(field_positions)) model_items_callback = operator.__annotations__["snapform_model"]["items"] with patch.object( self.addon.runtime, @@ -454,11 +476,6 @@ class AddonTest(unittest.TestCase): "NESTED_GIANT_LARGE_1_1", ], ) - self.assertIn("invoke_props_dialog", inspect.getsource(operator.invoke)) - self.assertNotIn( - '"DA3"', - inspect.getsource(operator), - ) def test_conversion_dialog_remembers_last_depth_model(self): operator = self.addon.ConvertImageToDepthPlane() @@ -466,6 +483,12 @@ class AddonTest(unittest.TestCase): is_property_set=lambda _name: False, ) context = SimpleNamespace( + object=SimpleNamespace( + data=SimpleNamespace(source="FILE", frame_duration=1) + ), + scene=SimpleNamespace( + snapform_settings=SimpleNamespace(video_frames=16) + ), window_manager=SimpleNamespace( invoke_props_dialog=lambda invoked, **_options: invoked, ), @@ -482,6 +505,85 @@ class AddonTest(unittest.TestCase): self.assertIs(result, operator) self.assertEqual(operator.snapform_model, "GIANT_1_1") + def test_remove_background_defaults_to_empty_image_frames(self): + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace(source="MOVIE", frame_duration=97), + image_user=SimpleNamespace(frame_duration=24), + ) + scene = SimpleNamespace( + snapform_settings=SimpleNamespace(video_frames=16) + ) + context = SimpleNamespace(object=source_object, scene=scene) + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + operator = self.addon.RemoveImageBackground() + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "input.png" + input_path.touch() + with patch.object( + self.addon.image_depth, + "prepare_image_input", + return_value=(input_path, False), + ): + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + self.assertEqual(scene.snapform_settings.video_frames, 24) + + def test_conversion_execute_uses_remove_background_frame_range(self): + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace( + name="clip.mp4", + source="MOVIE", + frame_duration=97, + size=(1920, 1080), + ), + image_user=SimpleNamespace(frame_duration=24), + ) + scene = SimpleNamespace( + snapform_settings=SimpleNamespace(video_frames=16) + ) + context = SimpleNamespace(object=source_object, scene=scene) + operator = self.addon.ConvertImageToDepthPlane() + operator.snapform_depth_source = "ESTIMATE" + operator.snapform_model = "BASE" + operator.snapform_use_ray_pose = True + operator.snapform_process_res_factor = 1.0 + operator.snapform_displacement_via = "MODIFIER" + operator.snapform_unshaded = False + operator.snapform_remove_background = False + operator.snapform_subdivision_level = 5 + operator.snapform_delete_original = False + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + run_job=lambda **options: calls.append(options) + or {"RUNNING_MODAL"}, + ), + ) + + with TemporaryDirectory() as directory: + input_path = Path(directory) / "clip.mp4" + input_path.touch() + with patch.object( + self.addon.image_depth, + "prepare_image_input", + return_value=(input_path, False), + ): + result = operator.execute(context) + + self.assertEqual(result, {"FINISHED"}) + self.assertEqual(scene.snapform_settings.video_frames, 24) + self.assertEqual(calls[0]["job_type"], "IMAGE_DEPTH") + def test_conversion_dialog_only_lists_installed_models(self): items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[ "snapform_model" @@ -494,21 +596,6 @@ class AddonTest(unittest.TestCase): items = items_callback(None, None) self.assertEqual([item[0] for item in items], ["SMALL", "BASE"]) - draw_source = inspect.getsource(self.addon.ConvertImageToDepthPlane.draw) - self.assertIn('text="No DA3 models installed"', draw_source) - self.assertIn('text="Download one in Host > Models"', draw_source) - self.assertNotIn("_dialog_requirements_ready", inspect.getsource(self.addon.ConvertImageToDepthPlane)) - - def test_depth_conversion_progress_is_provider_neutral(self): - starting_source = inspect.getsource( - self.addon.RunJob._starting_message - ) - finish_source = inspect.getsource(self.addon.RunJob._finish) - - self.assertIn('return "Starting depth estimation"', starting_source) - self.assertIn('"Depth plane is ready"', finish_source) - self.assertNotIn("DA3 image depth", starting_source) - self.assertNotIn("DA3 depth plane", finish_source) def test_update_status_formats_stage_progress(self): operator = self.addon.RunJob() @@ -550,12 +637,6 @@ class AddonTest(unittest.TestCase): "Estimating depth", ) - def test_depth_background_is_loaded_without_changing_source_empty(self): - source = inspect.getsource(self.addon.RunJob._build_depth_plane) - - self.assertIn("snapform_background_path=", source) - self.assertNotIn("_apply_removed_background", source) - def test_model_options_mark_only_models_that_are_not_downloaded(self): items_callback = self.addon.SnapFormSettings.__annotations__["model"][ "items" @@ -677,6 +758,71 @@ class AddonTest(unittest.TestCase): self.assertEqual(closed, [True]) self.assertIsNone(self.addon.operators.active_job()) + def test_temporary_frame_directory_is_removed_when_job_closes(self): + closed = [] + with TemporaryDirectory() as directory: + root = Path(directory) + temporary_directory = root / "inputs" / "frames" + temporary_directory.mkdir(parents=True) + (temporary_directory / "frame-0000.jpg").touch() + self.addon.operators._active_job = SimpleNamespace( + log_file=SimpleNamespace(close=lambda: closed.append(True)), + temporary_input_path=temporary_directory, + ) + + self.addon.operators.close_active_job() + + self.assertFalse(temporary_directory.exists()) + self.assertEqual(closed, [True]) + self.assertIsNone(self.addon.operators.active_job()) + + def test_prepare_image_input_passes_movie_file_through(self): + with TemporaryDirectory() as directory: + video_path = Path(directory) / "clip.mp4" + video_path.touch() + image = SimpleNamespace( + source="MOVIE", + filepath=str(video_path), + ) + + result, delete_input = self.addon.image_depth.prepare_image_input( + image, + object(), + ) + + self.assertEqual(result, video_path) + self.assertFalse(delete_input) + + def test_prepare_image_input_copies_sequence_frames(self): + with TemporaryDirectory() as directory: + root = Path(directory) + sequence_directory = root / "seq" + sequence_directory.mkdir() + (sequence_directory / "frame-0001.jpg").touch() + (sequence_directory / "frame-0002.jpg").touch() + image = SimpleNamespace( + source="SEQUENCE", + filepath=str(sequence_directory / "frame-0001.jpg"), + ) + temporary_directory = root / "temp-frames" + temporary_directory.mkdir() + with patch.object( + self.addon.runtime, + "temporary_input_directory", + return_value=temporary_directory, + ): + result, delete_input = ( + self.addon.image_depth.prepare_image_input( + image, + object(), + ) + ) + + self.assertEqual(result, temporary_directory) + self.assertTrue(delete_input) + self.assertTrue((temporary_directory / "frame.0001.jpg").is_file()) + self.assertTrue((temporary_directory / "frame.0002.jpg").is_file()) + def test_image_depth_model_and_ray_pose_are_independent(self): settings = SimpleNamespace( input_path="debug-input.png", @@ -853,17 +999,10 @@ class AddonTest(unittest.TestCase): str(runtime.DEFAULT_STORAGE_ROOT), ) - def test_ai_host_uses_authenticated_loopback_http(self): + def test_ai_host_uses_loopback_address(self): runtime = importlib.import_module("snapform.runtime") - request_source = inspect.getsource(runtime._host_http_request) - start_source = inspect.getsource(runtime.start_host) self.assertEqual(runtime.HOST_ADDRESS, "127.0.0.1") - self.assertIn('"X-SnapForm-Token"', request_source) - self.assertIn('"/health"', inspect.getsource(runtime._check_host)) - self.assertIn('with_name("ai_host.py")', start_source) - self.assertIn("ProxyHandler({})", inspect.getsource(runtime)) - self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path)) def test_closing_blender_terminates_host_it_started(self): runtime = importlib.import_module("snapform.runtime") @@ -890,13 +1029,6 @@ class AddonTest(unittest.TestCase): self.assertIsNone(runtime._host_connection) self.assertIsNone(runtime._registered_instance_id) - def test_host_jobs_complete_from_status_state(self): - source = inspect.getsource(self.addon.RunJob.modal) - - self.assertIn("if job.uses_host:", source) - self.assertIn('state == "succeeded"', source) - self.assertIn('{"failed", "cancelled"}', source) - def test_process_res_method_defaults_to_upper_bound_resize(self): process_res = self.addon.SnapFormSettings.__annotations__[ "process_res" @@ -959,6 +1091,202 @@ class AddonTest(unittest.TestCase): with self.subTest(parameter=parameter): self.assertIn(parameter, parameters) + def test_remove_background_request_passes_video_frames(self): + runtime = importlib.import_module("snapform.runtime") + settings = SimpleNamespace(device="AUTO", video_frames=24) + with patch.object( + runtime, + "runtime_python", + return_value=Path("python"), + ), patch.object( + runtime, + "model_directory", + return_value=Path("ben2_base"), + ): + request = runtime.remove_background_request( + settings, + Path("clip.mp4"), + Path("status.json"), + Path("result"), + "BEN2_BASE", + start_frame=70, + ) + + parameters = request.payload["parameters"] + self.assertEqual(parameters["video_frames"], 24) + self.assertEqual(parameters["start_frame"], 70) + + def test_extract_frames_request_passes_frame_range(self): + runtime = importlib.import_module("snapform.runtime") + settings = SimpleNamespace(device="AUTO", video_frames=24) + + request = runtime.extract_frames_request( + settings, + Path("clip.mp4"), + Path("status.json"), + Path("result"), + start_frame=70, + ) + + parameters = request.payload["parameters"] + self.assertEqual(request.payload["command"], "extract-frames") + self.assertEqual(parameters["input"], str(Path("clip.mp4"))) + self.assertEqual(parameters["output_dir"], str(Path("result"))) + self.assertEqual(parameters["video_frames"], 24) + self.assertEqual(parameters["start_frame"], 70) + + def test_new_job_paths_wraps_each_job_in_a_folder(self): + runtime = importlib.import_module("snapform.runtime") + with TemporaryDirectory() as directory: + jobs_root = Path(directory) + with patch.object( + runtime, + "jobs_directory", + return_value=jobs_root, + ): + paths = runtime.new_job_paths("RECONSTRUCT") + + self.assertTrue(paths["result"].is_dir()) + self.assertEqual(paths["result"].parent, jobs_root) + self.assertEqual(paths["log"].parent, paths["result"]) + self.assertEqual(paths["status"].parent, paths["result"]) + self.assertEqual(paths["status"].stem, paths["result"].name) + + def test_new_job_paths_uses_type_time_folder_name(self): + runtime = importlib.import_module("snapform.runtime") + with TemporaryDirectory() as directory: + jobs_root = Path(directory) + with patch.object( + runtime, + "jobs_directory", + return_value=jobs_root, + ): + depth_paths = runtime.new_job_paths( + self.addon.operators.IMAGE_DEPTH_JOB + ) + background_paths = runtime.new_job_paths( + self.addon.operators.IMAGE_BACKGROUND_JOB + ) + frames_paths = runtime.new_job_paths( + self.addon.operators.ORIGINAL_FRAMES_JOB + ) + + self.assertRegex( + depth_paths["result"].name, + r"^depth-plane_\d{8}-\d{6}$", + ) + self.assertTrue( + background_paths["result"].name.startswith( + "remove-background_" + ) + ) + self.assertTrue( + frames_paths["result"].name.startswith("depth-plane_") + ) + + def test_empty_image_start_frame_honors_image_user(self): + operators = self.addon.operators + + self.assertEqual( + operators.empty_image_start_frame( + SimpleNamespace( + data=SimpleNamespace(source="MOVIE", frame_duration=97), + image_user=SimpleNamespace( + frame_start=1, + frame_offset=70, + frame_duration=24, + ) + ) + ), + 71, + ) + self.assertEqual( + operators.empty_image_start_frame( + SimpleNamespace( + data=SimpleNamespace(source="MOVIE", frame_duration=97), + image_user=SimpleNamespace( + frame_start=10, + frame_offset=5, + frame_duration=24, + ) + ) + ), + 6, + ) + self.assertEqual( + operators.empty_image_start_frame( + SimpleNamespace(image_user=None) + ), + 1, + ) + + def test_image_depth_job_passes_empty_frame_offset(self): + settings = SimpleNamespace( + input_path="debug-input.png", + export_format="mini_npz", + model="BASE", + gs_model="GIANT_1_1", + process_res=504, + process_res_method="upper_bound_resize", + conf_thresh_percentile=40.0, + num_max_points=1_000_000, + show_cameras=True, + gs_views_interval=1, + video_frames=16, + device="AUTO", + use_ray_pose=False, + ref_view_strategy="saddle_balanced", + ) + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator.snapform_source_object_name = "Reference" + operator.snapform_model_key = "BASE" + operator.snapform_use_ray_pose = True + operator.snapform_process_res = 1024 + original_get = self.fake_bpy.data.objects.get + self.fake_bpy.data.objects.get = lambda _name: SimpleNamespace( + data=SimpleNamespace(source="MOVIE", frame_duration=97), + image_user=SimpleNamespace( + frame_start=1, + frame_offset=70, + frame_duration=24, + ), + ) + try: + with TemporaryDirectory() as directory: + input_path = Path(directory) / "input.png" + input_path.touch() + operator.snapform_input_path = str(input_path) + paths = { + "log": Path(directory) / "job.log", + "status": Path(directory) / "job.json", + "result": Path(directory) / "result", + } + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value=paths, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "reconstruct_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), + ) as reconstruct: + operator._prepare_command(None, settings) + finally: + self.fake_bpy.data.objects.get = original_get + + self.assertEqual(reconstruct.call_args.kwargs["start_frame"], 71) + def test_reconstruct_request_accepts_image_depth_overrides(self): runtime = importlib.import_module("snapform.runtime") settings = SimpleNamespace( @@ -1037,61 +1365,10 @@ class AddonTest(unittest.TestCase): self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname) self.assertIn("DEFAULT_CLOSED", panel.bl_options) - def test_debug_panel_uses_one_export_button_and_dynamic_settings(self): - debug_source = inspect.getsource(self.addon.DebugPanel.draw) - da3_source = inspect.getsource(self.addon.Da3DebugPanel.draw) - ben2_source = inspect.getsource(self.addon.Ben2DebugPanel.draw) - - self.assertIn('layout.prop(settings, "device")', debug_source) - self.assertEqual(da3_source.count("export_row.operator("), 1) - self.assertIn('settings.export_format == "glb"', da3_source) - self.assertIn('settings.export_format == "gs_ply"', da3_source) - self.assertIn('"gs_model" if settings.export_format', da3_source) - self.assertIn("layout.prop(settings, model_property)", da3_source) - self.assertNotIn("model_row", da3_source) - self.assertIn("and model_ready(selected_model_key(settings))", da3_source) - self.assertNotIn('text="Ready"', da3_source) - self.assertNotIn('text="Missing"', da3_source) - self.assertNotIn("License:", da3_source) - self.assertIn('operator.job_type = "RECONSTRUCT"', da3_source) - self.assertIn( - 'layout.prop(settings, "background_input_path")', - ben2_source, - ) - self.assertIn('text="Remove Background"', ben2_source) - self.assertIn("and model_ready(BEN2_MODEL.key)", ben2_source) - self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', ben2_source) - def test_host_panel_exposes_shared_host_state_and_controls(self): self.assertIn(self.addon.HostPanel, self.addon.CLASSES) self.assertIn(self.addon.ControlHost, self.addon.CLASSES) - source = inspect.getsource(self.addon.HostPanel.draw) - self.assertEqual(self.addon.HostPanel.bl_label, "Host") - self.assertIn('text="Runtime is not installed"', source) - self.assertIn('text="Install Runtime"', source) - self.assertIn("runtime.host_status()", source) - self.assertIn('f"Running (Clients:', source) - self.assertIn('else "Running (Busy)"', source) - self.assertIn('"READY": "RADIOBUT_ON"', source) - self.assertIn('"BUSY": "TIME"', source) - self.assertNotIn('"READY": "PLAY"', source) - self.assertNotIn('text=f"Task:', source) - self.assertIn('elif state == "STOPPED":', source) - self.assertNotIn("Blender Clients:", source) - self.assertNotIn("status['port']", source) - self.assertIn('icon="TRASH"', source) - self.assertIn('icon="FILE_REFRESH"', source) - self.assertIn('icon="PAUSE"', source) - self.assertNotIn('text="Open Log"', source) - preferences_source = inspect.getsource(self.addon.SnapFormPreferences.draw) - self.assertIn('text="Open Host Log"', preferences_source) - - def test_host_start_control_does_not_wait_for_model_runtime(self): - source = inspect.getsource(self.addon.runtime.enable_host) - - self.assertIn("start_host()", source) - self.assertNotIn("ensure_host()", source) def test_ai_operation_restarts_a_manually_stopped_host(self): runtime = self.addon.runtime @@ -1160,40 +1437,6 @@ class AddonTest(unittest.TestCase): runtime._host_auto_start = previous_auto_start runtime._host_snapshot = previous_snapshot - def test_models_panel_is_collapsed_and_assigns_each_download_target(self): - draw_source = inspect.getsource( - self.addon.ModelsPanel.draw - ) - model_source = inspect.getsource( - self.addon.ModelsPanel._draw_model - ) - - self.assertIn("for model in MODELS.values()", draw_source) - self.assertIn( - "self._draw_model(layout, BEN2_MODEL, running)", - draw_source, - ) - self.assertIn("operator.model_key = model.key", model_source) - self.assertIn( - "download_row.enabled = runtime_ready() and not running", - model_source, - ) - self.assertNotIn('label(text="Not downloaded"', model_source) - self.assertNotIn("Isolated Runtime ready", draw_source) - self.assertNotIn("License:", model_source) - self.assertEqual( - self.addon.ModelsPanel.bl_label, - "Models", - ) - self.assertEqual( - self.addon.ModelsPanel.bl_parent_id, - self.addon.HostPanel.bl_idname, - ) - self.assertEqual( - self.addon.ModelsPanel.bl_options, - {"DEFAULT_CLOSED"}, - ) - def test_ben2_debug_job_uses_its_image_and_model(self): settings = SimpleNamespace( background_input_path="input.png", @@ -1285,6 +1528,91 @@ class AddonTest(unittest.TestCase): self.assertEqual(prepared_paths, paths) self.assertEqual(remove_background.call_args.args[1], input_path) + def test_image_background_job_accepts_video_input(self): + settings = SimpleNamespace(device="AUTO", video_frames=16) + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB + + with TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "clip.mp4" + input_path.touch() + operator.snapform_input_path = str(input_path) + paths = { + "log": root / "job.log", + "status": root / "job.json", + "result": root / "result", + } + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value=paths, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "remove_background_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), + ) as remove_background: + command, prepared_paths = operator._prepare_command( + None, + settings, + ) + + self.assertIsInstance(command, self.addon.runtime.HostJobRequest) + self.assertEqual(remove_background.call_args.args[1], input_path) + + def test_original_frames_job_uses_frame_extraction_request(self): + settings = SimpleNamespace(device="AUTO", video_frames=16) + operator = self.addon.RunJob() + operator.job_type = self.addon.operators.ORIGINAL_FRAMES_JOB + + with TemporaryDirectory() as directory: + root = Path(directory) + input_path = root / "clip.mp4" + input_path.touch() + operator.snapform_input_path = str(input_path) + paths = { + "log": root / "job.log", + "status": root / "job.json", + "result": root / "result", + } + with patch.object( + self.addon.runtime, + "new_job_paths", + return_value=paths, + ), patch.object( + self.addon.runtime, + "runtime_ready", + return_value=True, + ), patch.object( + self.addon.runtime, + "extract_frames_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), + ) as extract_frames: + command, prepared_paths = operator._prepare_command( + None, + settings, + ) + + self.assertIsInstance(command, self.addon.runtime.HostJobRequest) + self.assertEqual(prepared_paths, paths) + self.assertEqual(extract_frames.call_args.args[1], input_path) + self.assertEqual( + extract_frames.call_args.kwargs["start_frame"], + 1, + ) + def test_completed_image_background_job_applies_result_with_undo(self): calls = [] self.fake_bpy.ops = SimpleNamespace( @@ -1323,6 +1651,41 @@ class AddonTest(unittest.TestCase): ], ) + def test_completed_image_background_job_applies_foreground_sequence(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + apply_removed_background=lambda **options: calls.append( + ("apply", options) + ) + or {"FINISHED"}, + ), + ed=SimpleNamespace( + undo_push=lambda **options: calls.append(("undo", options)) + or {"FINISHED"}, + ), + ) + operator = self.addon.RunJob() + with TemporaryDirectory() as directory: + result_directory = Path(directory) + (result_directory / "foreground.0000.png").touch() + (result_directory / "foreground.0001.png").touch() + job = SimpleNamespace( + result_directory=result_directory, + source_object_name="Reference", + ) + + operator._apply_removed_background(job) + + self.assertEqual( + calls[0][1]["snapform_output_path"], + str(result_directory), + ) + self.assertEqual( + calls[1], + ("undo", {"message": "Remove Background"}), + ) + def test_completed_original_background_job_builds_with_one_undo(self): calls = [] self.fake_bpy.ops = SimpleNamespace( @@ -1339,6 +1702,7 @@ class AddonTest(unittest.TestCase): ) operator = self.addon.RunJob() job = SimpleNamespace( + kind=self.addon.operators.ORIGINAL_BACKGROUND_JOB, result_directory=Path("result"), source_object_name="Reference", subdivision_level=5, @@ -1360,6 +1724,53 @@ class AddonTest(unittest.TestCase): ("undo", {"message": "Convert to Depth Plane"}), ) + def test_completed_original_frames_job_builds_with_color_frames(self): + calls = [] + self.fake_bpy.ops = SimpleNamespace( + snapform=SimpleNamespace( + build_original_depth_plane=lambda **options: calls.append( + ("build", options) + ) + or {"FINISHED"}, + ), + ed=SimpleNamespace( + undo_push=lambda **options: calls.append(("undo", options)) + or {"FINISHED"}, + ), + ) + operator = self.addon.RunJob() + with TemporaryDirectory() as directory: + result_directory = Path(directory) + (result_directory / "frame.0001.jpg").touch() + (result_directory / "frame.0002.jpg").touch() + job = SimpleNamespace( + kind=self.addon.operators.ORIGINAL_FRAMES_JOB, + result_directory=result_directory, + source_object_name="Reference", + subdivision_level=5, + delete_source=False, + displacement_via="MODIFIER", + unshaded=False, + ) + + operator._build_original_depth_plane(job) + + build_options = calls[0][1] + self.assertEqual(build_options["snapform_background_path"], "") + self.assertEqual( + build_options["snapform_color_frames_dir"], + str(result_directory), + ) + self.assertEqual( + build_options["snapform_depth_frames_dir"], + str(result_directory), + ) + self.assertFalse(build_options["snapform_delete_source"]) + self.assertEqual( + calls[1], + ("undo", {"message": "Convert to Depth Plane"}), + ) + def test_ben2_download_command_filters_large_alternative_weights(self): runtime = importlib.import_module("snapform.runtime") with patch.object( @@ -1414,12 +1825,34 @@ class AddonTest(unittest.TestCase): any("bootstrap_runtime.py" in str(value) for value in command) ) - def test_preferences_install_button_is_disabled_while_busy(self): - draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw) + def test_install_runtime_command_passes_configured_pypi_index(self): + runtime = importlib.import_module("snapform.runtime") + with patch.object( + runtime, + "uv_executable", + return_value="C:/uv/uv.exe", + ), patch.object( + runtime, + "preferences", + return_value=SimpleNamespace( + torch_backend="CU128", + pypi_index_url="https://pypi.tuna.tsinghua.edu.cn/simple", + ), + ), patch.object( + runtime, + "runtime_directory", + return_value=Path("runtime"), + ), patch.object( + runtime, + "models_directory", + return_value=Path("models"), + ): + command = runtime.install_runtime_command(Path("status.json")) - self.assertIn( - "install_row.enabled = active_job() is None", - draw_source, + index_position = command.index("--pypi-index") + self.assertEqual( + command[index_position + 1], + "https://pypi.tuna.tsinghua.edu.cn/simple", ) def test_completed_export_opens_job_directory(self): @@ -1435,6 +1868,34 @@ class AddonTest(unittest.TestCase): self.assertEqual(len(calls), 1) self.assertIn("filepath", calls[0]) + def test_completed_job_removes_status_file(self): + operator = self.addon.RunJob() + with TemporaryDirectory() as directory: + status_path = Path(directory) / "status.json" + status_path.write_text( + '{"state": "succeeded", "progress": 1.0, "message": "done"}', + encoding="utf-8", + ) + job = SimpleNamespace( + kind="INSTALL_RUNTIME", + cancelled=False, + log_file=None, + status_path=status_path, + ) + context = SimpleNamespace( + scene=SimpleNamespace( + snapform_settings=SimpleNamespace( + status="", + progress=0.0, + ), + ), + ) + with patch.object(operator, "_redraw_ui"): + result = operator._finish(context, job, 0) + + self.assertEqual(result, {"FINISHED"}) + self.assertFalse(status_path.exists()) + def test_status_bar_draws_progress_and_cancel_button(self): progress_calls = [] operator_calls = [] @@ -1478,68 +1939,6 @@ class AddonTest(unittest.TestCase): [("snapform.cancel_job", {"text": "", "icon": "X"})], ) - def test_image_depth_status_uses_same_progress_layout(self): - progress_calls = [] - operator_calls = [] - - class Layout: - def row(self, align=False): - return self - - def label(self, **_options): - return None - - def progress(self, **options): - progress_calls.append(options) - - def operator(self, operator, **options): - operator_calls.append((operator, options)) - - self.addon.operators._active_job = SimpleNamespace( - kind=self.addon.operators.IMAGE_DEPTH_JOB, - ) - context = SimpleNamespace( - scene=SimpleNamespace( - snapform_settings=SimpleNamespace( - status="Estimating depth", - progress=0.25, - ), - ), - ) - - self.addon.draw_status_bar( - SimpleNamespace(layout=Layout()), - context, - ) - - self.assertEqual( - progress_calls, - [ - { - "factor": 0.25, - "type": "BAR", - "text": "Estimating depth", - } - ], - ) - self.assertEqual( - operator_calls, - [("snapform.cancel_job", {"text": "", "icon": "X"})], - ) - - def test_progress_is_appended_before_cancel_button(self): - register_source = inspect.getsource(self.addon.register) - draw_source = inspect.getsource(self.addon.draw_status_bar) - - self.assertIn("STATUSBAR_HT_header.append", register_source) - self.assertNotIn("STATUSBAR_HT_header.prepend", register_source) - self.assertLess( - draw_source.index("progress_row.progress("), - draw_source.index('"snapform.cancel_job"'), - ) - self.assertIn("progress_row.ui_units_x = 12.0", draw_source) - self.assertNotIn('label(text="SnapForm"', draw_source) - def test_ui_redraw_targets_status_bar_and_view_3d_areas(self): redrawn = [] status_regions = ( diff --git a/tests/test_ai_host.py b/tests/test_ai_host.py index e74669a..dc18646 100644 --- a/tests/test_ai_host.py +++ b/tests/test_ai_host.py @@ -1,5 +1,4 @@ import importlib.util -import inspect import sys import tempfile import unittest @@ -112,21 +111,6 @@ class HostTest(unittest.TestCase): self.assertEqual(loads, ["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): fake_models = SimpleNamespace( status=lambda: {"ben2": "", "da3": ""}, @@ -218,6 +202,51 @@ class HostTest(unittest.TestCase): self.host.remove_discovery(path, "owner") 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__": unittest.main() diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py index b50d3e0..6a2fdf0 100644 --- a/tests/test_bootstrap_runtime.py +++ b/tests/test_bootstrap_runtime.py @@ -1,7 +1,9 @@ import importlib.util +import io import json import os import subprocess +import tarfile import tempfile import time import unittest @@ -25,12 +27,24 @@ def load_bootstrap_runtime(): 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): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" 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.write_bytes(b"uv") args = SimpleNamespace( @@ -42,7 +56,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ) as mocked_run, patch.object( @@ -64,13 +78,7 @@ class BootstrapRuntimeTest(unittest.TestCase): f"addict=={bootstrap.ADDICT_VERSION}", da3_command, ) - self.assertTrue( - any( - "https://github.com/ByteDance-Seed/Depth-Anything-3/" - "archive/test-commit.tar.gz" in str(value) - for value in da3_command - ) - ) + self.assertIn(f"depth-anything-3 @ {da3_uri}", da3_command) self.assertFalse( any("git+" in str(value) for value in da3_command) ) @@ -93,13 +101,7 @@ class BootstrapRuntimeTest(unittest.TestCase): for command in commands if any("ben2 @" in str(value) for value in command) ) - self.assertTrue( - any( - "https://github.com/PramaLLC/BEN2/archive/" - "test-ben2-commit.tar.gz" in str(value) - for value in ben2_command - ) - ) + self.assertIn(f"ben2 @ {ben2_uri}", ben2_command) self.assertFalse( any("git+" in str(value) for value in ben2_command) ) @@ -162,6 +164,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -181,7 +184,7 @@ class BootstrapRuntimeTest(unittest.TestCase): bootstrap.platform, "machine", return_value="arm64", - ), patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( + ), patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run", ), patch.object( @@ -329,6 +332,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -364,7 +368,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ) return [] - with patch.object(bootstrap, "preflight", return_value=({}, False)), patch.object( + with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( bootstrap, "run_stage", side_effect=fake_run_stage, @@ -425,7 +429,8 @@ class BootstrapRuntimeTest(unittest.TestCase): "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] self.assertEqual(first.kwargs["stage"], 2) self.assertEqual(first.kwargs["stages"], 8) @@ -525,146 +530,79 @@ class BootstrapRuntimeTest(unittest.TestCase): self.assertEqual(first.kwargs["stages"], 8) 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): bootstrap = load_bootstrap_runtime() self.assertEqual(bootstrap.format_elapsed(12), "12s") self.assertEqual(bootstrap.format_elapsed(83), "1m 23s") - def test_preflight_rejects_low_disk_space(self): + def test_resolve_indexes_uses_configured_pypi_index(self): bootstrap = load_bootstrap_runtime() - args = SimpleNamespace( - backend="cu128", - runtime_dir=str(Path.cwd()), + bootstrap.resolve_indexes( + SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple") ) - with patch.object( - bootstrap.shutil, - "disk_usage", - return_value=SimpleNamespace(free=1024**3), - ), patch.object(bootstrap, "reachable", return_value=True): - with self.assertRaisesRegex(RuntimeError, "disk space"): - bootstrap.preflight(args) - - def test_preflight_rejects_unreachable_network(self): - bootstrap = load_bootstrap_runtime() - args = SimpleNamespace( - backend="cu128", - runtime_dir=str(Path.cwd()), - ) - with patch.object( - bootstrap.shutil, - "disk_usage", - return_value=SimpleNamespace(free=100 * 1024**3), - ), patch.object(bootstrap, "reachable", return_value=False): - with self.assertRaisesRegex(RuntimeError, "Cannot reach"): - bootstrap.preflight(args) - - def test_preflight_passes_when_checks_ok(self): - bootstrap = load_bootstrap_runtime() - args = SimpleNamespace( - backend="cu128", - runtime_dir=str(Path.cwd()), - ) - with patch.object( - bootstrap.shutil, - "disk_usage", - return_value=SimpleNamespace(free=100 * 1024**3), - ), patch.object(bootstrap, "reachable", return_value=True): - resolved, used_mirror = bootstrap.preflight(args) - - self.assertFalse(used_mirror) self.assertEqual( - resolved["pypi"], - "https://pypi.org/simple/", - ) - - def test_preflight_falls_back_to_mirrors(self): - bootstrap = load_bootstrap_runtime() - args = SimpleNamespace( - backend="cu128", - runtime_dir=str(Path.cwd()), - ) - - def fake_reachable(url, timeout=bootstrap.PREFLIGHT_TIMEOUT): - if "github.com" in url: - return True - if "tuna" in url or "hf-mirror" in url or "aliyun" in url: - return True - return False - - with patch.object( - bootstrap.shutil, - "disk_usage", - return_value=SimpleNamespace(free=100 * 1024**3), - ), patch.object( - bootstrap, - "reachable", - side_effect=fake_reachable, - ): - resolved, used_mirror = bootstrap.preflight(args) - - self.assertTrue(used_mirror) - self.assertEqual( - resolved["pypi"], + bootstrap._resolved_pypi_index, "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() - args = SimpleNamespace( - backend="cu128", - runtime_dir=str(Path.cwd()), - ) with patch.dict( os.environ, - { - "PIP_INDEX_URL": "https://pypi.example/simple", - "HF_ENDPOINT": "https://hf.example", - "SNAPFORM_TORCH_INDEX": "https://torch.example/whl/cu128", - }, - ), patch.object( - bootstrap.shutil, - "disk_usage", - return_value=SimpleNamespace(free=100 * 1024**3), - ), patch.object(bootstrap, "reachable", return_value=True): - resolved, used_mirror = bootstrap.preflight(args) - - self.assertFalse(used_mirror) - self.assertEqual(resolved["pypi"], "https://pypi.example/simple") - self.assertEqual(resolved["hf"], "https://hf.example") + {"PIP_INDEX_URL": "https://pypi.example/simple"}, + ): + bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) self.assertEqual( - resolved["torch"], - "https://torch.example/whl/cu128", + bootstrap._resolved_pypi_index, + "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() previous_pypi = bootstrap._resolved_pypi_index - previous_hf = bootstrap._resolved_hf_endpoint bootstrap._resolved_pypi_index = ( "https://pypi.tuna.tsinghua.edu.cn/simple" ) - bootstrap._resolved_hf_endpoint = "https://hf-mirror.com" try: env = bootstrap.child_environment() finally: bootstrap._resolved_pypi_index = previous_pypi - bootstrap._resolved_hf_endpoint = previous_hf self.assertEqual( env["PIP_INDEX_URL"], "https://pypi.tuna.tsinghua.edu.cn/simple", ) - self.assertEqual(env["HF_ENDPOINT"], "https://hf-mirror.com") def test_child_environment_injects_registry_proxy(self): bootstrap = load_bootstrap_runtime() @@ -700,6 +638,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) uv_path = root / "uv" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -711,7 +650,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ), patch.object( @@ -737,6 +676,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) (runtime_dir / "partial-python.exe").write_text("junk") args = SimpleNamespace( uv="uv", @@ -747,7 +687,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ), patch.object( @@ -765,6 +705,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir(parents=True) + self.create_bundled_sources(root) (runtime_dir / "Scripts").mkdir() (runtime_dir / "Scripts" / "python.exe").write_text("python") (runtime_dir / "snapform_ready.json").write_text( @@ -781,7 +722,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ), patch.object( @@ -803,6 +744,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) args = SimpleNamespace( uv="uv", runtime_dir=str(runtime_dir), @@ -812,7 +754,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", side_effect=fake_run, @@ -826,12 +768,13 @@ class BootstrapRuntimeTest(unittest.TestCase): self.assertEqual(status["progress"], 1.0) 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() with tempfile.TemporaryDirectory() as directory: root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) uv_path = root / "uv" / "uv.exe" args = SimpleNamespace( uv=str(uv_path), @@ -842,7 +785,7 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ), patch.object( @@ -855,15 +798,101 @@ class BootstrapRuntimeTest(unittest.TestCase): ), patch.object( bootstrap.shutil, "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) - mocked_download.assert_called_once() - call_args = mocked_download.call_args - self.assertEqual(call_args.args[0], uv_path) - self.assertEqual(call_args.args[1], str(root / "status.json")) - self.assertEqual(call_args.kwargs["stage"], 1) - self.assertEqual(call_args.kwargs["stage_label"], "uv") + self.assertEqual(uv_path.read_bytes(), b"uv-binary") + all_commands = run_commands + install_commands + self.assertFalse( + any( + "github.com" in str(value) + for command in all_commands + for value in command + ) + ) + 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): bootstrap = load_bootstrap_runtime() @@ -871,6 +900,7 @@ class BootstrapRuntimeTest(unittest.TestCase): root = Path(directory) runtime_dir = root / "runtime" runtime_dir.mkdir() + self.create_bundled_sources(root) uv_path = root / "uv.exe" uv_path.write_bytes(b"uv") args = SimpleNamespace( @@ -882,17 +912,15 @@ class BootstrapRuntimeTest(unittest.TestCase): ben2_commit="test-ben2-commit", 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, "run", ), patch.object( bootstrap, "run_install", - ), patch.object(bootstrap, "download_uv") as mocked_download: + ): bootstrap.install(args) - mocked_download.assert_not_called() - def test_extract_uv_unpacks_standalone_binary(self): bootstrap = load_bootstrap_runtime() with tempfile.TemporaryDirectory() as directory: @@ -906,6 +934,112 @@ class BootstrapRuntimeTest(unittest.TestCase): 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__": unittest.main() diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index 2783a24..b5e8bc7 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -143,49 +143,29 @@ class DownloadProgressTest(unittest.TestCase): self.assertFalse(status_path.exists()) - def test_resolve_hf_endpoint_keeps_direct_endpoint(self): - with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, - "reachable", - return_value=True, - ): - endpoint, used_mirror = self.worker.resolve_hf_endpoint() + def test_download_endpoints_keeps_direct_endpoint(self): + with patch.dict(os.environ, {"HF_ENDPOINT": ""}): + endpoints = self.worker.download_endpoints() - self.assertEqual(endpoint, "https://huggingface.co/") - self.assertFalse(used_mirror) + self.assertEqual( + endpoints, + ["https://huggingface.co"], + ) - def test_resolve_hf_endpoint_falls_back_to_mirror(self): - with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, - "reachable", - side_effect=lambda url, timeout=None: "hf-mirror" in url, - ): - endpoint, used_mirror = self.worker.resolve_hf_endpoint() + def test_download_endpoints_respects_environment_override(self): + with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}): + endpoints = self.worker.download_endpoints() - self.assertEqual(endpoint, "https://hf-mirror.com") - self.assertTrue(used_mirror) + self.assertEqual( + endpoints, + ["https://hf.example"], + ) - def test_resolve_hf_endpoint_respects_environment_override(self): - with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}), patch.object( - self.worker, - "reachable", - ) as mocked_reachable: - endpoint, used_mirror = self.worker.resolve_hf_endpoint() + def test_download_endpoints_strips_trailing_slash(self): + with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}): + endpoints = self.worker.download_endpoints() - self.assertEqual(endpoint, "https://hf.example") - self.assertFalse(used_mirror) - mocked_reachable.assert_not_called() - - def test_resolve_hf_endpoint_keeps_direct_when_mirrors_unreachable(self): - with patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, - "reachable", - return_value=False, - ): - endpoint, used_mirror = self.worker.resolve_hf_endpoint() - - self.assertEqual(endpoint, "https://huggingface.co/") - self.assertFalse(used_mirror) + self.assertEqual(endpoints, ["https://hf.example"]) if __name__ == "__main__": diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 2b9ba92..9edd366 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -1,5 +1,4 @@ import importlib.util -import inspect import sys import tempfile import unittest @@ -60,6 +59,22 @@ class ImageDepthTest(unittest.TestCase): 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): with tempfile.TemporaryDirectory() as directory: npz_path = Path(directory) / "results.npz" @@ -72,6 +87,261 @@ class ImageDepthTest(unittest.TestCase): 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): with tempfile.TemporaryDirectory() as directory: npz_path = Path(directory) / "results.npz" @@ -205,16 +475,6 @@ class ImageDepthTest(unittest.TestCase): (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): source = SimpleNamespace( data=SimpleNamespace(size=(800, 400)), @@ -233,37 +493,6 @@ class ImageDepthTest(unittest.TestCase): ): 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): image = SimpleNamespace(size=(1080, 1920)) @@ -272,11 +501,6 @@ class ImageDepthTest(unittest.TestCase): 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): source_object = SimpleNamespace( name="obj_img", @@ -292,17 +516,20 @@ class ImageDepthTest(unittest.TestCase): self.image_depth.depth_data_name(source_object), "img.001.depth", ) - estimate_source = inspect.getsource( - self.image_depth.build_depth_plane - ) - original_source = inspect.getsource( - self.image_depth.build_image_depth_plane - ) - self.assertIn( - "depth_name = depth_name or depth_data_name(source_object)", - estimate_source, - ) - self.assertIn("depth_data_name(source_object)", original_source) + + def test_copy_depth_image_renames_still_source(self): + class FakeImage: + def __init__(self, name): + self.name = name + + def copy(self): + return FakeImage("img.png.001") + + source_object = SimpleNamespace(data=FakeImage("img.png")) + + 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): source_object = SimpleNamespace( @@ -311,213 +538,14 @@ class ImageDepthTest(unittest.TestCase): self.assertEqual( self.image_depth.removed_background_data_name(source_object), - "img.rb", + "img.foreground", ) source_object.data.name = "img.png.001" self.assertEqual( 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): class Material: surface_displacement_method = "BUMP" @@ -530,12 +558,6 @@ class ImageDepthTest(unittest.TestCase): "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): source = SimpleNamespace( data=SimpleNamespace(size=(800, 400)), diff --git a/tests/test_pack.py b/tests/test_pack.py index 2376663..7cd3f98 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -1,4 +1,7 @@ +import ast import importlib.util +import io +import tarfile import tempfile import unittest from pathlib import Path @@ -19,25 +22,148 @@ def load_pack(): class PackTest(unittest.TestCase): - def test_builds_extension_root_layout(self): + def test_builds_platform_specific_extension_zips(self): pack = load_pack() with tempfile.TemporaryDirectory() as directory: original_dist = pack.DIST_DIR pack.DIST_DIR = Path(directory) try: - output_path = pack.build_extension() + output_paths = pack.build_all() finally: pack.DIST_DIR = original_dist - self.assertEqual(output_path.name, "SnapForm.v0.9.32.zip") - with ZipFile(output_path) as archive: - names = set(archive.namelist()) + self.assertEqual( + {path.name for path in output_paths}, + { + "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("blender_manifest.toml", names) - self.assertIn("image_depth.py", names) - self.assertNotIn("snapform_depth_plane.py", names) - self.assertIn("worker.py", names) - self.assertIn("ai_host.py", names) - self.assertNotIn("importer.py", names) - self.assertNotIn("snapform/__init__.py", names) + self.assertIn("__init__.py", names) + self.assertIn("blender_manifest.toml", names) + self.assertIn("image_depth.py", names) + self.assertIn("worker.py", names) + self.assertIn("ai_host.py", names) + self.assertNotIn("snapform/__init__.py", names) + self.assertNotIn("importer.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() diff --git a/tests/test_worker.py b/tests/test_worker.py index 766c4ac..a203a4d 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -1,5 +1,4 @@ import importlib.util -import inspect import json import os import sys @@ -42,27 +41,11 @@ class WorkerTest(unittest.TestCase): 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): with self.assertRaises(self.worker.TaskCancelled): self.worker.check_cancelled(lambda: True) 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): mps_torch = SimpleNamespace( @@ -115,13 +98,7 @@ class WorkerTest(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "MPS was requested"): self.worker.select_device("mps", cpu_only_torch) - def test_optional_background_removal_runs_before_da3_model_load(self): - source = inspect.getsource(self.worker.reconstruct) - - self.assertLess( - source.index("infer_foreground("), - source.index("DepthAnything3.from_pretrained"), - ) + def test_reconstruct_parser_accepts_background_model_directory(self): args = self.worker.parser().parse_args( [ "reconstruct", @@ -175,6 +152,199 @@ class WorkerTest(unittest.TestCase): self.assertEqual(images, [image_path]) 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): with tempfile.TemporaryDirectory() as directory: root = Path(directory) @@ -235,10 +405,6 @@ class WorkerTest(unittest.TestCase): self.worker, "create_download_progress_class", return_value=object, - ), patch.object( - self.worker, - "resolve_hf_endpoint", - return_value=(self.worker.HF_ENDPOINT, False), ): args.model_dir = str(Path(directory) / "model") args.status = str(Path(directory) / "status.json") @@ -248,13 +414,17 @@ class WorkerTest(unittest.TestCase): calls[0]["allow_patterns"], ["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 = [] 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( repository="PramaLLC/BEN2", model_dir="model", @@ -269,23 +439,151 @@ class WorkerTest(unittest.TestCase): self.worker, "create_download_progress_class", return_value=object, - ), patch.object( - self.worker, - "resolve_hf_endpoint", - return_value=("https://hf-mirror.com", True), ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( self.worker, - "write_status", - ) as mocked_write: + "download_from_modelscope", + ) 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.status = str(Path(directory) / "status.json") self.worker.download_model(args) - self.assertEqual( - os.environ["HF_ENDPOINT"], - "https://hf-mirror.com", - ) - messages = [call.args[2] for call in mocked_write.call_args_list] - self.assertTrue(any("(mirror)" in message for message in messages)) + + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT) + self.assertEqual(mocked_modelscope.call_count, 1) + + 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): inference_options = [] @@ -368,6 +666,166 @@ class WorkerTest(unittest.TestCase): ) 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): inference_options = [] export_calls = [] @@ -500,6 +958,95 @@ class WorkerTest(unittest.TestCase): infer_gs_values = [options["infer_gs"] for _, options in inference_options] 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__": unittest.main() diff --git a/uv.lock b/uv.lock index 116ff06..4d59669 100644 --- a/uv.lock +++ b/uv.lock @@ -4,5 +4,5 @@ requires-python = ">=3.11" [[package]] name = "snapform" -version = "0.9.30" +version = "0.10.10" source = { virtual = "." }