From c5b77c5c855cf4b93ec7049ccd7a961b6c7c211b Mon Sep 17 00:00:00 2001 From: nan Date: Mon, 17 Aug 2026 10:37:55 +0800 Subject: [PATCH] Add ONNX model inference implementations for MoGe-2, BEN2, and DA3 - Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing. - Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing. - Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy. - Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`. - Developed unit tests for server functionalities in `test_server.py`, including model caching and job management. - Enhanced job server to handle frame extraction without requiring models. --- AGENTS.md | 5 +- README.md | 15 +- docs/blendjob.md | 330 +- docs/depth-plane.md | 2 +- docs/development.md | 28 +- docs/features.md | 34 +- docs/index.md | 2 +- docs/server.md | 115 + docs/service.md | 104 - docs/surface.md | 22 +- src/snapform/__init__.py | 39 +- src/snapform/ai_service/__init__.py | 1 - src/snapform/ai_service/jobs.py | 2277 -------------- src/snapform/ai_service/server.py | 145 - src/snapform/blender/depth_plane.py | 388 +-- src/snapform/blender/surface.py | 129 +- src/snapform/blender_manifest.toml | 2 +- src/snapform/blendjob/__init__.py | 38 +- src/snapform/blendjob/client.py | 60 +- .../blendjob/{service.py => controller.py} | 61 +- src/snapform/blendjob/environment.py | 351 +++ .../blendjob/installer/environment.py | 13 + src/snapform/blendjob/operator.py | 260 +- src/snapform/blendjob/runner.py | 71 + src/snapform/blendjob/runtime.py | 596 ++++ src/snapform/blendjob/server.py | 172 +- src/snapform/client.py | 521 ---- src/snapform/env_revision.py | 4 - src/snapform/install_env.py | 921 ------ src/snapform/job_runtime.py | 54 + src/snapform/operators.py | 1277 ++++---- src/snapform/panel.py | 85 +- src/snapform/preferences.py | 21 +- src/snapform/properties.py | 152 +- src/snapform/server/__init__.py | 64 + src/snapform/server/background.py | 170 + src/snapform/server/da3.py | 233 ++ src/snapform/server/geometry.py | 616 ++++ src/snapform/server/media.py | 306 ++ .../{models.py => server/model_catalog.py} | 136 +- src/snapform/server/model_download.py | 444 +++ .../{ai_service => server}/model_manager.py | 50 +- src/snapform/server/moge2.py | 304 ++ .../ben2_onnx.py => server/onnx_ben2.py} | 1 + .../da3_onnx.py => server/onnx_da3.py} | 1 + .../moge_onnx.py => server/onnx_moge2.py} | 1 + .../{ai_service => server}/onnx_runtime.py | 0 src/snapform/status_bar.py | 55 - tests/test_addon.py | 2739 +++-------------- tests/test_ben2_onnx.py | 8 +- tests/test_da3_onnx.py | 8 +- tests/test_download_progress.py | 212 +- tests/test_install_env.py | 1405 +-------- tests/test_models.py | 2 +- tests/test_moge_onnx.py | 4 +- tests/test_pack.py | 19 +- tests/{test_service.py => test_server.py} | 196 +- tests/test_worker.py | 502 ++- 58 files changed, 5840 insertions(+), 9931 deletions(-) create mode 100644 docs/server.md delete mode 100644 docs/service.md delete mode 100644 src/snapform/ai_service/__init__.py delete mode 100644 src/snapform/ai_service/jobs.py delete mode 100644 src/snapform/ai_service/server.py rename src/snapform/blendjob/{service.py => controller.py} (75%) create mode 100644 src/snapform/blendjob/environment.py create mode 100644 src/snapform/blendjob/installer/environment.py create mode 100644 src/snapform/blendjob/runner.py create mode 100644 src/snapform/blendjob/runtime.py delete mode 100644 src/snapform/client.py delete mode 100644 src/snapform/env_revision.py delete mode 100644 src/snapform/install_env.py create mode 100644 src/snapform/job_runtime.py create mode 100644 src/snapform/server/__init__.py create mode 100644 src/snapform/server/background.py create mode 100644 src/snapform/server/da3.py create mode 100644 src/snapform/server/geometry.py create mode 100644 src/snapform/server/media.py rename src/snapform/{models.py => server/model_catalog.py} (62%) create mode 100644 src/snapform/server/model_download.py rename src/snapform/{ai_service => server}/model_manager.py (73%) create mode 100644 src/snapform/server/moge2.py rename src/snapform/{ai_service/ben2_onnx.py => server/onnx_ben2.py} (99%) rename src/snapform/{ai_service/da3_onnx.py => server/onnx_da3.py} (99%) rename src/snapform/{ai_service/moge_onnx.py => server/onnx_moge2.py} (99%) rename src/snapform/{ai_service => server}/onnx_runtime.py (100%) delete mode 100644 src/snapform/status_bar.py rename tests/{test_service.py => test_server.py} (50%) diff --git a/AGENTS.md b/AGENTS.md index 31e5c2a..269dd81 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,7 +5,7 @@ ## 项目文档 - 动手前先读 `docs/index.md`，它是全部文档的入口，按主题链接并帮助理解项目结构与代码逻辑 -- 主题映射：功能与使用 `docs/features.md`、Depth Plane `docs/depth-plane.md`、Surface `docs/surface.md`、输出格式 `docs/outputs.md`、AI Service `docs/service.md`、开发与许可证 `docs/development.md` +- 主题映射：功能与使用 `docs/features.md`、Depth Plane `docs/depth-plane.md`、Surface `docs/surface.md`、输出格式 `docs/outputs.md`、Job Server `docs/server.md`、开发与许可证 `docs/development.md` - 修改代码或执行方式后，同步更新对应主题文档与 `README.md`，保持文档与实现一致 ## Markdown 规范 @@ -35,4 +35,5 @@ - Operator、Menu、Panel、AddonPreferences... Class 用大驼峰命名，不加任何前缀。Operator 动宾结构，比如`RenderImage`。Menu、Panel 以它们本身为后缀，比如`RenderPanel` - 根据项目名提炼项目缩写，作为 bl_idname 和 自定义 Property 的前缀以区分其他功能和属性，以`Omoo`为例 - bl_idname 用下划线小写命名，Operator 的 bl_idname 以`omoo.*`为前缀，比如`omoo.render_panel`；Menu 的以`OMOO_MT_*`为前缀；Panel 的以`OMOO_PT_*`为前缀。比如`OMOO_PT_render_panel` - - 自定义 Property 用下划线小写命名，以`omoo_*`为前缀。比如`omoo_progress_factor` \ No newline at end of file + - Operator 自己声明的 Property 使用无项目前缀的业务名称，比如`input_path`、`resolution_level` + - 只有直接挂载到公共 Blender RNA 类型的自定义 Property 使用项目前缀，比如`Scene.omoo_settings` diff --git a/README.md b/README.md index 830227b..cd2433c 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ SnapForm 是一个 Blender 扩展，可以把一张照片、一段视频或一 扩展不内置 uv 或 AI 运行库。首次使用会先下载并校验固定版本的 uv，随后自动安装轻量的 ONNX Runtime 环境，再下载 `MoGe-2 ViT-S Normal` 和 `BEN2 Base`；不安装 PyTorch，也不下载各模型的 Python 源码包。其他 AI 模型可在 `Models` panel 中按需下载，请保持网络畅通，首次安装需要一些时间。 -安装时 uv、ONNX Environment、MoGe-2 和 BEN2 各自使用独立的 0–100% 进度；下载阶段显示速度与连续字节进度。 +安装时 Environment、MoGe-2 和 BEN2 各自使用独立的 0–100% 进度；Environment 这一轮内部依次完成 uv、Python 环境和 packages。进入后续模型时进度归零，但不显示 `2/2` 等阶段编号。模型源发生瞬时错误时会切换到另一来源继续下载。 ## 快速上手 @@ -64,7 +64,7 @@ MoGe-2 ViT-S、ViT-B、ViT-L、DA3 Small、DA3 Base 与 BEN2 全部使用 ONNX 以及分割线下方的 `Keep Original`。 界面会显示计算所得的实际 Process Res。剔除和重拓扑参数直接在 Geometry Nodes 修改器中调整；`Unshaded` 可使用与 Depth Plane 相同的 Emission 材质。普通材质的 -Bump Scale 会在材质组内乘以 `0.001`，并按透视表面的深度空间缩放。Surface attributes 统一使用 `o_` 前缀；Service 准备过程显示 +Bump Scale 会在材质组内乘以 `0.001`，并按透视表面的深度空间缩放。Surface attributes 统一使用 `o_` 前缀；Server 准备过程显示 `Preparing surface data n/6`，可看到当前阶段与总步骤。 ### 视频与序列帧 @@ -75,19 +75,20 @@ Session 逐帧推理，显存峰值不随视频帧数平方增长；Debug 多帧 ## 任务结果保存在哪里 -任务结果保存在 `~/.snapform/jobs/` 下。Job Server 接受任务时自动生成 opaque Job ID，并创建同名文件夹；Blender 不预先申请 ID 或目录。Packed Image 和序列帧的临时输入由 Blender 在任务结束后清理，Service 产物直接写入该 Job 目录。视频或序列帧生成的面片会引用其中的序列帧，请保留对应文件夹，不要随意删除。 +任务结果保存在 Runtime Storage Root 的 `jobs/` 下，默认是 `~/.snapform/jobs/`。Job Server 接受任务时自动生成 opaque Job ID，并创建同名文件夹；Blender 不预先申请 ID 或目录。Packed Image 和序列帧的临时输入由 Blender 在任务结束后清理，Server 产物直接写入该 Job 目录。视频或序列帧生成的面片会引用其中的序列帧，请保留对应文件夹，不要随意删除。 ## 常见问题 - 首次使用需要下载模型，如果下载慢或失败，稍后在 `Models` panel 里重试即可 -- 安装 Environment 时提示下载超时：新版会自动探测并切换到镜像、失败时自动重试一次；如果仍然 - 超时，在扩展偏好设置的 `PyPI Index` 中填入镜像地址（如清华）后重新安装即可 +- 安装 Environment 时 uv 下载失败会自动尝试 OmooLab 镜像；Python packages 需要指定镜像时设置 + `UV_INDEX_URL` 后重新安装 +- 默认模型下载会在 Hugging Face 与 OmooLab 镜像之间自动回退；某一来源暂时断线不会阻止后续模型尝试另一来源 - 找不到右键入口时，请确认选中了 Empty，且显示方式为 Image - 动画面片不动：检查 Empty 的 `Image Frames` 是否大于 1，并确认播放场景帧 ## 技术文档 -AI、帧提取与后续模型下载统一提交给当前 Blender 专属的 Job Server。DA3、MoGe-2 与 Debug 的 Request 和 Job Type 分别一一对应，不使用混合 Backend Request。BlendJob 接受任务时创建无类型前缀的 opaque-ID 工作目录，并通过 `JobContext` 交给执行逻辑。运行中进度通过 HTTP 轮询，已加载的 ONNX Session 由 Server 的 Models Resource 统一查询与清理。Environment 安装仍由不依赖 Service 的 Bootstrap 流程完成，不创建 Job 目录。 +AI、帧提取与后续模型下载统一提交给当前 Blender 专属的 Job Server。`runtime = BlenderJobRuntime(...)` 独占 Storage Root、Environment、Server 进程、Status Bar 和当前 Job；Python 与依赖通过 `environment` 字典声明，BlendJob 负责完整安装。JobServer 使用单 Worker FIFO Queue，忙碌时接受并顺序执行后续任务。内建 Status Bar 使用兼容 Blender 4.5 回调所有权机制的稳定函数对象。SnapForm 的 `post_install(runtime)` 使用 `runtime.request()` 下载默认模型。`server/__init__.py` 集中注册全部公开 Job；业务实现按 `media.py`、`geometry.py`、`background.py`、`moge2.py` 与 `da3.py` 直接拆分。Job-based Operator 继承 `runtime.JobOperatorBase`，直接返回参数字典，并在可选的 `response(context, result)` 中完成 Blender 后处理。模型 catalog、下载和 Session 分别由 `model_catalog.py`、`model_download.py` 与 `model_manager` Resource 管理，不存在循环 import。模型下载固定尝试 Hugging Face，再尝试 R2，两者失败即结束当前 Job。BlendJob 接受任务时在统一 Storage Root 下创建 opaque-ID 工作目录。 详细功能与技术内容见 [docs/index.md](docs/index.md)。 -可复用的本地 Job Service API 与插件接入方式见 [docs/blendjob.md](docs/blendjob.md)。 +可复用的本地 Job Server API 与插件接入方式见 [docs/blendjob.md](docs/blendjob.md)。 diff --git a/docs/blendjob.md b/docs/blendjob.md index 636b442..da62f9a 100644 --- a/docs/blendjob.md +++ b/docs/blendjob.md @@ -1,179 +1,255 @@ # BlendJob 开发者指南 -## 命名与边界 +## 目标与边界 -通用 module 命名为 `blendjob`。这个名字表达的是 Blender 插件与独立 Job 进程之间的桥接，不绑定 AI、模型框架或 SnapForm。 +`blendjob` 连接 Blender Operator 与独立 Python Job Server。开发者新增任务时只需要一个 Blender Operator 和一个 Server Handler，不需要编写提交线程、HTTP 轮询、取消、进度条、Server 进程或 Environment 安装 Operator。 -- `JobClient`：运行在 Blender Python 中，通过 HTTP 提交、查询和取消 Job -- `JobServer`：对 FastAPI 的通用封装，负责 Job 注册、单任务调度、状态与取消 -- `JobContext`：单个 Job 的运行上下文，提供工作目录、进度、取消与 Resource -- `ModalJobOperatorBase`：可直接继承的 Blender Modal Operator，内置提交线程、轮询、进度、取消和终态处理 -- `Resource`：由 Job Server 持有、跨多个 Job 复用的状态对象 -- `SnapForm Job Server`：SnapForm 创建并注册专属 Job 后得到的服务实例；代码变量使用 `server` 或 `snapform_server` +- `BlenderJobRuntime`：Blender 侧唯一运行时，持有 Storage Root、Environment、Server 进程和当前 Job +- `JobOperatorBase`：绑定 Runtime 后供业务 Operator 直接继承的 Modal Mixin +- `JobServer`：独立 Python 进程中的 Job 注册、状态、取消、工作目录和 Resource +- `JobContext`：Handler 的工作目录、Storage Root、进度、取消与 Resource 接口 +- `JobResult`：Server 成功结果及其 Job 文件 +- `ServerController`：Runtime 内部的专属 Server 进程控制器 -`blendjob` 不知道模型、输出文件或 Blender Operator。插件专属代码负责注册 Job 和 Resource，并把参数交给自己的执行逻辑。 +## 项目配置 -## 最小接入 - -服务端只创建 Server 并注册业务函数： +Server 实例直接定义在项目的 `server` package： ```python -server = JobServer("Example", jobs_directory="D:/example/jobs") +# snapform/server/__init__.py + +from blendjob import JobServer -@server.job("generate-depth") -def generate_depth(context, parameters): - output = context.directory / "depth.npz" - run_depth(parameters["input"], output) - context.update(1.0, "Depth is ready", output=str(output)) +server = JobServer("SnapForm Job Server") ``` -Blender 端直接继承一个基类。普通任务只需要声明 Controller、Job Type 和怎么传值： +Blender 侧把这个实例交给 Runtime。实例不会跨进程传输；Runtime 从导出该实例的 package 自动找到独立 Python 入口。 ```python -class GenerateDepth(ModalJobOperatorBase): - bl_idname = "example.generate_depth" - bl_label = "Generate Depth" - controller_factory = staticmethod(service_controller) - job_type = "generate-depth" +# snapform/job_runtime.py - def job_parameters(self, context): - return {"input": context.scene.source_path} +from .server import server + + +environment = { + "python": "3.12", + "packages": [ + "numpy==2.4.2", + "pillow==12.1.1", + "av==17.0.0", + "trimesh==4.11.2", + "huggingface-hub==1.4.1", + ], + "platform_packages": { + "windows": ["onnxruntime-directml==1.24.4"], + "default": ["onnxruntime==1.24.4"], + }, +} + + +def post_install(runtime): + runtime.request( + "download-model", + {"model": "MOGE2_VITS_NORMAL"}, + ) + + +runtime = BlenderJobRuntime( + server, + storage_root=storage_root, + environment=environment, + post_install=post_install, +) + +JobOperatorBase = runtime.JobOperatorBase ``` -这个 Operator 已经具备异步提交、Server 生成 Job ID 与目录、HTTP 轮询、进度、Esc 取消、错误报告和 Timer 清理。只有需要把产物导回 Blender 时，才覆盖 `on_job_succeeded(context, job, status)`；产物目录是 `job.directory`。 +Runtime 按约定管理 `.venv`、Manifest、安装状态与日志路径。BlendJob 根据 `environment` 字典安装固定 Python、通用的 FastAPI/Uvicorn Server 依赖、项目 packages 以及当前平台 packages。配置会生成稳定 Hash；Python 或依赖声明变化后 Environment 自动判定为需要重装，不再需要项目维护 `install_env.py` 或手工 Revision。 -## Resource +`platform_packages` 优先选择 `windows`、`macos` 或 `linux`，没有对应项时使用 `default`。基础 packages 与平台 packages 会合并安装。 -Job 表示一次有终态的操作；Resource 表示在 Service 生命周期内被多个 Job 共用的状态。Resource 可以按需实现 `snapshot()`、`clear()` 和 `close()`： +`post_install` 是可选的普通 callable，在 Environment 可用且 Server 能启动后执行。它不限定为 Job 列表，可以执行任意 Python；SnapForm 用它发起默认模型下载。连续 request 会分别重置进度，Status Bar 不为未知数量的 post-install 操作显示 `2/2` 等阶段总数。post-install 失败不会撤销已经安装成功的 Environment。 + +`runtime.register()` 自动注册 Environment 安装、Job 取消、Server 启动、停止、重启和日志 Operator，并挂载内建 Status Bar。Status Bar 使用可由 Blender 4.5 标记所有权的稳定函数回调，Runtime 卸载时移除同一个回调对象。业务 `CLASSES` 只包含业务类型。 ```python -models = server.add_resource("models", ModelManager(models_directory)) +CLASSES = ( + GenerateMoGe2Depth, +) -@server.job("generate-depth") -def generate_depth(context, parameters): - models = context.resource("models") - session = models.get(parameters["model"], parameters["device"]) +def register(): + for class_type in CLASSES: + bpy.utils.register_class(class_type) + runtime.register() + + +def unregister(): + runtime.unregister() + for class_type in reversed(CLASSES): + bpy.utils.unregister_class(class_type) ``` -`JobServer` 在同一把调度锁下阻止运行中清理 Resource，关闭时按注册逆序调用 `close()`。Resource Snapshot 自动进入 `/health` 的 `resources` 字段。 +启用 Runtime 不会立即启动 Server。第一次提交 Job 或显式调用 Start Server 时才启动当前 Blender 专属进程。 -## 服务端使用 +## Blender Operator -开发者先创建 `JobServer`，再用装饰器注册 Job 类型： +最小 Operator 只声明 `job_type` 与自己的 Blender Property： ```python -from blendjob import JobCancelled, JobServer +class GenerateMoGe2Depth(JobOperatorBase, bpy.types.Operator): + bl_idname = "snapform.generate_moge2_depth" + bl_label = "Generate MoGe-2 Depth" + job_type = "generate-moge2-depth" - -server = JobServer("Example Job Server", jobs_directory="D:/example/jobs") - - -@server.job("generate-depth") -def generate_depth(context, parameters): - input_path = parameters["input"] - output_path = context.directory / "depth.npz" - context.update(0.1, "Loading image") - context.check_cancelled() - - run_model(input_path, output_path) - - context.update( - 1.0, - "Depth is ready", - state="succeeded", - output=output_path, + input_path: bpy.props.StringProperty() + resolution_level: bpy.props.IntProperty( + default=5, + min=0, + max=9, ) ``` -处理函数签名固定为 `handler(context, parameters)`。`parameters` 必须可由 JSON 表示。客户端提交的 `JobRequest` 不包含 ID；`JobServer` 接受请求时自动生成 opaque Job ID，并创建 `//`。处理函数通过 `context.job_id` 与 `context.directory` 获取它们。客户端不创建 Job 目录，也不应把输出目录作为业务参数传入。Job API 只保存内存状态和少量结果元数据，不创建 `status.json`。 +执行 `bpy.ops.snapform.generate_moge2_depth()` 调用的就是这个 Operator。`JobOperatorBase` 提供它的 `execute()` 和 `modal()`，不会再创建或调用隐藏的 Job Operator。 -`context.update()` 支持以下通用字段： +没有覆盖 `request()` 时，Operator 声明的全部 Property 会作为参数提交： -- `progress`：`0.0–1.0` -- `message`：面向界面的当前步骤 -- `state`：通常自动推导，也可显式设置 `queued`、`running`、`succeeded`、`failed`、`cancelling` 或 `cancelled` -- `error`：失败信息 -- 其它关键字：作为插件自定义状态字段原样返回 - -长耗时流程应在推理前后和循环内调用 `context.check_cancelled()`。底层库不可中断时，取消会在它返回后的最近检查点生效。 - -## Blender 端底层 API - -`JobClient` 只依赖 Python 标准库，适合直接放入 Blender 插件： - -```python -from blendjob import JobClient, JobRequest, ServiceController - - -client = JobClient("127.0.0.1", port) -request = JobRequest.create( - "generate-depth", - input="D:/images/source.png", -) -submitted = client.submit_request(request) -job_id = submitted["job_id"] -job_directory = submitted["directory"] - -status = client.status(job_id) -if status["state"] == "running": - update_blender_progress(status["progress"], status["message"]) +```json +{ + "input_path": "D:/image.png", + "resolution_level": 5 +} ``` -专属 Service 进程由 `ServiceController` 管理。插件只提供启动命令、工作目录和日志路径，端口、Instance ID、进程句柄、连接验证以及 `start()`、`ensure()`、`enable()`、`restart()`、`stop()`、`detach()` 都由 Controller 负责： +后端可直接忽略不使用的字段。BlendJob 不提供 Property 排除表。 + +### request + +需要读取 Blender Context 或生成临时输入时覆盖： ```python -controller = ServiceController( - lambda port, instance_id: [ - str(environment_python), - "server.py", - "--port", str(port), - "--instance-id", instance_id, - "--parent-pid", str(blender_pid), - ], - cwd_factory=lambda: storage_root, - log_path_factory=lambda: storage_root / "service.log", -) - -connection = controller.ensure() -submitted = connection.client.submit_request(request) -controller.stop() +def request(self, context): + input_path, temporary = prepare_image_input( + context.object.data, + context.scene, + ) + self._temporary_input = input_path if temporary else None + return { + "input": str(input_path), + "resolution_level": self.resolution_level, + } ``` -对应 HTTP API： +返回值是完整参数，纯替换默认 Property 参数，不做合并。它必须是可由 JSON 表示的 `dict`。 -- `GET /health`：查询服务身份、忙闲状态和插件扩展状态 -- `POST /jobs`：提交 Job -- `GET /jobs/{job_id}`：轮询进度与终态 -- `DELETE /jobs/{job_id}`：请求取消 -- `GET /resources`：查询全部 Resource -- `GET /resources/{resource_name}`：查询一个 Resource -- `POST /resources/{resource_name}/clear`：空闲时清理 Resource -- `POST /shutdown`：关闭当前专属服务 +### response -Service 只监听 `127.0.0.1`，专属实例由启动它的 `ServiceController` 保存随机端口与 Instance ID。这个本地专属进程协议不额外使用 Token 鉴权。 +Server 成功后，Runtime 在 Blender 主线程调用： -`ModalJobOperatorBase` 默认使用 Blender Window Manager 的进度接口。插件若有自己的状态栏，可覆盖 `update_job_ui()` 与 `redraw_job_ui()`；这些属于界面定制，不是接入 Job 所必需的代码。运行态使用通用 `ModalJobState`，只包含 Controller、Job ID、目录、进度和取消状态。插件特有的完成参数可由 `create_job_data()` 返回，并通过 `job.data` 读取，不应加入通用状态类。 +```python +def response(self, context, result): + build_moge2_depth_plane( + context, + source_object=bpy.data.objects[self._source_object_name], + npz_path=result.file("depth"), + ) +``` -## 生命周期建议 +`response()` 是可选的。未覆盖时 Job 成功即返回 `FINISHED`。开发者可以直接调用 Python 函数，也可以调用其它 Blender Operator；BlendJob 不限制后处理方式。 -每个 Blender 实例为每个插件启动一个专属 Job Server。随机端口与 Instance ID 由 `ServiceController` 生成并只保存在内存中，不写公共发现文件，也不扫描其它本地服务。 +`cleanup()` 同样可选，在成功、失败和取消后都会执行，适合删除临时输入。 -启动服务时传入 Blender 的 PID。服务同时监听显式 `/shutdown` 和父进程存活状态，因此正常禁用插件时可立即停止，Blender 异常退出后也能自行退出。这个方案允许同一机器出现多个服务进程，换取明确的所有权和插件隔离。 +业务代码直接返回 `dict`。BlendJob 没有公开 `JobRequest` 包装类型；HTTP Payload 由 Runtime 内部构造。 -模型与 Environment 可以继续放在插件自己的共享 Storage Root 中；共享磁盘缓存不代表共享服务进程。若多个 Blender 可能同时写同一个模型目录，插件仍需对下载或更新操作单独加文件锁。 +## 直接 Request -## 在其它插件复用 +不需要 Blender Operator 时，可以直接向 Runtime 发起一个同步 request： -当前实现位于 `src/snapform/blendjob/`，没有 SnapForm 依赖。复用时复制整个 `blendjob` package，并在插件专属入口中完成以下工作： +```python +def downloaded(result): + print(result.directory) -1. 创建带 `jobs_directory` 的 `JobServer` -2. 用 `@server.job()` 注册插件自己的处理函数 -3. 用 `server.add_resource()` 注册需要跨 Job 复用的状态 -4. 用安装环境中的 Uvicorn 运行 `server.create_app(instance_id)` -5. Blender 端用 `ServiceController` 管理启动、停止、重启和连接检查 -6. Blender Operator 继承 `ModalJobOperatorBase`，声明 `controller_factory`、`job_type` 与 `job_parameters()` -7. 需要导入产物时覆盖 `on_job_succeeded()` -后续若多个项目稳定复用，可把 `blendjob` 发布为独立 Python package；现阶段保持源码内嵌可以避免为 Blender 扩展增加安装依赖。 +result = runtime.request( + "download-model", + {"model": "MOGE2_VITS_NORMAL"}, + response=downloaded, +) +``` + +`runtime.request(job_type, parameters, response=None)` 负责启动 Server、提交、轮询、失败处理与构造 `JobResult`，返回值就是 Server response。可选的 `response(result)` 在发起 request 的线程执行。它是同步 API，适合 `post_install`、后台线程或脚本；Blender 主线程中的交互功能仍应使用 `JobOperatorBase`，避免阻塞界面。 + +## Server Handler + +Handler 直接注册为 `handler(context, parameters)`： + +```python +@server.job("generate-moge2-depth") +def generate_moge2_depth(context, parameters): + model = context.resource("model_manager").get( + parameters["model"], + parameters["device"], + ) + context.progress(0.1, "Loading MoGe-2") + context.check_cancelled() + + depth = model.infer( + parameters["input"], + resolution_level=parameters["resolution_level"], + ) + output = context.directory / "results.npz" + save_depth(output, depth) + return {"depth": output.name} +``` + +返回值保存在终态的 `result` 字段。Blender 端的 `JobResult.file("depth")` 会把相对路径限制在当前 Job Directory 内并验证文件存在。 + +`JobContext` 提供： + +- `storage_root`：Runtime 传入的唯一共享根目录 +- `directory`：`storage_root/jobs/` +- `progress(value, message, **details)`：发布运行进度 +- `check_cancelled()`：在安全检查点响应取消 +- `resource(name)`：访问 Server 生命周期 Resource + +Handler 不需要也不应管理 `queued`、`succeeded`、`failed` 等状态；这些状态由 `JobServer` 内部更新，因此公开 API 不再同时提供语义重复的 `update()`。Environment 安装发生在 Server 启动前，仍通过 Runtime 的 Bootstrap 状态文件回传进度，不属于 `JobContext`。 + +## Storage Root + +Runtime 是 Storage Root 的唯一配置者。`JobServer` 不需要在项目代码中声明 `jobs_directory`，启动后固定使用： + +```text +/ +├── .venv/ +├── environment.json +├── jobs/ +├── models/ +├── tools/ +├── environment.log +├── environment-status.json +├── server.log +└── manifest.json +``` + +需要 Storage Root 的 Resource 使用延迟工厂： + +```python +@server.resource("model_manager") +def create_model_manager(job_server): + return ModelManager(job_server.storage_root / "models") +``` + +Resource 在独立进程绑定 Storage Root 后创建，关闭 Server 时按注册逆序关闭。 + +模型等领域数据应由 Server Resource 管理。Blender 可以通过 `runtime.resource("model_manager")` 查询 Resource Snapshot，不应自行拼接模型目录或检查模型文件。 + +## Server 生命周期 + +每个 Blender 实例为当前插件持有一个专属 Server。Runtime 自动处理随机端口、Instance ID、父 Blender PID、启动日志、连接验证、重启和停止。 + +Server 只监听 `127.0.0.1`。当前本地专属进程协议不使用一次性 Token，不发现或复用其它 Blender 启动的 Server。 + +`JobServer` 内建一个标准库单 Worker FIFO Queue。新 Job 会立即得到 ID 和目录，空闲时执行，忙碌时保持 `queued`；同一时刻只运行一个 Handler。排队任务可以取消，调用方不需要自己等待 Server Idle，`runtime.request()` 也不使用额外的 Idle 轮询补丁。 + +模型和 Environment 可以共享同一 Storage Root；共享磁盘目录不表示共享 Server 进程。多个进程可能同时下载同一模型时，业务层仍需要文件锁。 diff --git a/docs/depth-plane.md b/docs/depth-plane.md index a7fa719..08f3646 100644 --- a/docs/depth-plane.md +++ b/docs/depth-plane.md @@ -18,7 +18,7 @@ 参数会集中显示在独立边框中。`Remove Background` 与 `Unshaded` 直接显示，不再使用 Options 折叠； `Keep Original` 默认关闭，并在分割线下方保持为最后一项。仅当开启 `Remove Background` 时才需要 `BEN2 Base`，并把 `.foreground` 的 Alpha 作为独立 - Mask，不在 Service 中改写深度值。保存 `.foreground` 时保留 BEN2 Alpha，但 RGB 使用未经修改的原始输入； + Mask，不在 Server 中改写深度值。保存 `.foreground` 时保留 BEN2 Alpha，但 RGB 使用未经修改的原始输入； 因此材质直接用 `.foreground` 同时提供颜色和 Mask，不再额外加载原图 Image Texture。 Image Empty 引用视频或序列帧时，DA3 与 MoGe-2 会按 Empty 的 Image Frames 数值逐帧估计 深度（未设置时为完整帧数，显式设为 1 时只处理 1 帧），起始帧跟随 Empty 的 Offset，与背景蒙版范围一致； diff --git a/docs/development.md b/docs/development.md index ae32aa8..f041a2e 100644 --- a/docs/development.md +++ b/docs/development.md @@ -15,28 +15,36 @@ uv run python pack.py ├── src/ │ └── snapform/ │ ├── blender_manifest.toml -│ ├── install_env.py -│ ├── models.py │ ├── operators.py │ ├── panel.py │ ├── preferences.py │ ├── properties.py -│ ├── client.py +│ ├── job_runtime.py │ ├── blendjob/ │ │ ├── client.py +│ │ ├── controller.py +│ │ ├── environment.py +│ │ ├── installer/environment.py │ │ ├── operator.py -│ │ ├── service.py +│ │ ├── runner.py +│ │ ├── runtime.py │ │ └── server.py │ ├── blender/ │ │ ├── depth_plane.py │ │ └── surface.py -│ └── ai_service/ -│ ├── server.py +│ └── server/ +│ ├── __init__.py +│ ├── model_catalog.py +│ ├── model_download.py │ ├── model_manager.py -│ ├── jobs.py -│ ├── ben2_onnx.py -│ ├── da3_onnx.py -│ ├── moge_onnx.py +│ ├── media.py +│ ├── geometry.py +│ ├── background.py +│ ├── moge2.py +│ ├── da3.py +│ ├── onnx_ben2.py +│ ├── onnx_da3.py +│ ├── onnx_moge2.py │ └── onnx_runtime.py └── tests/ ``` diff --git a/docs/features.md b/docs/features.md index 628113d..6044d63 100644 --- a/docs/features.md +++ b/docs/features.md @@ -12,12 +12,12 @@ ## 当前功能 - 使用 Blender Extension manifest，可通过 `Install from Disk` 安装 -- 在独立 Environment 中通过常驻 Service 运行 ONNX Runtime；MoGe-2、DA3 与 BEN2 的模型执行及后处理完全不使用 PyTorch，避免污染 Blender Python 进程 -- Service 的深度模型缓存只保留 DA3 或 MoGe-2 其中一个，切换时主动释放前一个；BEN2 使用独立缓存， +- 在独立 Environment 中通过常驻 Server 运行 ONNX Runtime；MoGe-2、DA3 与 BEN2 的模型执行及后处理完全不使用 PyTorch，避免污染 Blender Python 进程 +- Server 的深度模型缓存只保留 DA3 或 MoGe-2 其中一个，切换时主动释放前一个；BEN2 使用独立缓存， 可与当前深度模型共存 -- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Service；请求带随机 Token -- 常驻 Service 复用 ONNX Runtime session、BEN2 与最近使用的深度模型，连续执行时不再重复加载 -- 多个 Blender 实例自动连接同一个共享 Service；`AI Service` 面板显示状态、端口、客户端与模型缓存，并提供安装、启动、清理、重启和停止控制 +- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Server；连接由随机端口与 Instance ID 归属当前 Runtime +- 常驻 Server 复用 ONNX Runtime session、BEN2 与最近使用的深度模型，连续执行时不再重复加载 +- 每个 Blender 实例持有自己的 Job Server；`Job Server` 面板显示状态与模型缓存，并提供安装、启动、清理、重启和停止控制 - Environment、模型和任务结果默认保存在 `~/.snapform`，可由多个 Blender 版本共享 - 当前仅启用 FP32 DA3 Small 和 Base；它们不随首次安装下载，Giant 1.1 与 Nested Giant Large 1.1 暂停使用 - 支持 BEN2 Base 单图背景剔除测试，输出透明 PNG 与推理耗时 @@ -45,7 +45,7 @@ - 转换对话框选择推理模式、处理分辨率与背景移除；细分和置换方式在 Geometry Nodes 修改器中调整 - Debug 的 DA3 输出或 BEN2 背景剔除完成后只打开对应 job 文件夹，不修改当前场景 - 完整保留每次任务生成的所有官方文件 -- 安装和模型下载使用后台子进程；推理提交给常驻 Service，状态栏实时显示阶段、百分比进度条与取消按钮 +- 安装和模型下载使用后台子进程；推理提交给常驻 Server，状态栏实时显示阶段、百分比进度条与取消按钮 - 后台任务运行时，Install、Download 和 Export 按钮保持原位置并变灰，任务完成后恢复 ## 使用 @@ -57,22 +57,20 @@ 3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能 （Image Empty 右键菜单）使用的推理设备，默认取本机最快的 DirectML/CoreML provider，否则 CPU，且只列出 本机支持的设备。 -4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn； - uv 由 Bootstrap 下载并校验，DA3、BEN2 与 MoGe 不安装源码包， - 并在 Environment 安装后自动下载 `MoGe-2 ViT-S Normal` 和 FP16 `BEN2 Base` 推荐模型。Environment 安装 - 首次安装依次显示 `uv`、`ONNX Environment`、`MoGe-2 ViT-S Normal` 与 `BEN2 Base`；每个阶段 - 都独立从 0 走到 100%，进入下一阶段时进度条归零。uv 已存在时省略其阶段并重新编号。 +4. 在 `Job Server` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn； + uv 由 BlendJob Environment 安装器下载并校验，DA3、BEN2 与 MoGe 不安装源码包， + 并在 Environment 安装后自动下载 `MoGe-2 ViT-S Normal` 和 FP16 `BEN2 Base` 推荐模型。首次安装 + 依次显示 `Environment`、`MoGe-2 ViT-S Normal` 与 `BEN2 Base`；每项都独立从 0 走到 100%， + 进入下一项时进度条归零，模型 post-install 不显示 `2/2` 等阶段编号。uv、Python Environment 与 packages 合并在第一轮， + 两个默认模型由普通的 `post_install(runtime)` 发起正式 Server Job。 安装完成后直接写入 Environment Manifest，不再额外导入全部组件进行验证；`DA3 Base` 等其他模型可稍后在 `Models` 面板按需下载。 - 状态文字使用 `阶段/总数 当前操作`，例如 - `3/4 Downloading MoGe-2 ViT-S Normal 8.4 MB/s`。依赖安装只显示阶段与已解析的 - 包总数；模型下载显示速度，例如 + 依赖安装显示当前安装输出；模型下载显示速度，例如 `Downloading DA3 Small 14.2 MB/s`；镜像来源显示为 `Downloading DA3 Small (mirror) 14.2 MB/s`。下载按字节连续推进，重试或切换镜像时进度不会倒退。 Environment 安装阶段与 `Models` 面板单独下载使用相同的模型名称格式，不会改写为通用的 `Downloading packages`。 - Environment 安装后即使模型下载中断也不影响 Service 使用，可稍后在 `Models` 面板单独补下。 - 安装不做启动网络预检：依赖安装默认依次尝试清华、阿里镜像并以 pypi.org 兜底；模型下载由 - Downloader 自行探测回退到 OmooLab 镜像，切换结果写入 `runtime.json` 的 `models_from_r2` 字段， - 后续下载直接走 R2；需要手动指定依赖镜像时可通过 + Environment 安装后即使模型下载中断也不影响 Server 使用，可稍后在 `Models` 面板单独补下。 + 安装不做启动网络预检；每个模型都先尝试 Hugging Face，失败后尝试 OmooLab R2，两者都失败时当前 Job 失败。 + 下载源不写入持久状态；需要手动指定依赖镜像时可通过 `UV_INDEX_URL` 环境变量覆盖。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 如需其他模型，展开 `Models` panel，找到模型并点击该行的 `Download`。 diff --git a/docs/index.md b/docs/index.md index 14072d4..5b17a5e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -6,6 +6,6 @@ - [Depth Plane](depth-plane.md)：深度面片的数据、节点与材质流程 - [Surface](surface.md)：透视网格、Geometry Nodes 后处理与材质法线 - [输出格式](outputs.md)：Debug 输出类型、参数和任务产物 -- [AI Service](service.md)：Environment、模型缓存、任务目录与运行机制 +- [Job Server](server.md)：Environment、模型缓存、任务目录与运行机制 - [BlendJob 开发者指南](blendjob.md)：可复用 Job Server、Blender Client 与装饰器注册 API - [开发与许可证](development.md)：开发验证、打包、上游版本与许可证 diff --git a/docs/server.md b/docs/server.md new file mode 100644 index 0000000..c3e78c8 --- /dev/null +++ b/docs/server.md @@ -0,0 +1,115 @@ +# Job Server + +## 运行架构 + +SnapForm 在独立的 Python 3.12 Environment 中运行 HTTP Job Server，Blender 只通过标准库 HTTP Client 提交和轮询任务，不把计算依赖载入 Blender Python。Server 只监听 `127.0.0.1`，由 BlendJob `ServerController` 保存随机端口与 Instance ID；专属本地进程协议不使用额外 Token 鉴权。 + +每个 Blender 实例启动并持有一个 SnapForm 专属 Server，不发现或复用其它 Blender、其它插件启动的本地 Server。多个 Blender 可以对应多个进程，但模型、Environment 与任务产物仍可共享同一个 Storage Root。 + +Environment、模型和任务结果默认保存在 `~/.snapform`： + +```text +~/.snapform/ +├── .venv/ +├── environment.json +├── models/ +├── jobs/ +├── manifest.json +├── environment.log +├── environment-status.json +└── server.log +``` + +不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。连接端口与 Instance ID 只保存在启动它们的 Blender 进程内存中。 + +通用 HTTP 层位于 `blendjob` module。`JobServer` 包装 FastAPI 并提供 Job、FIFO Queue 与 Resource API，`JobClient` 是无第三方依赖的 Blender 侧访问器。SnapForm 的 `server/__init__.py` 直接导出 `server` 实例、集中注册全部公开 Job，并把 `ModelManager` 注册为延迟创建的 `model_manager` Resource。完整接入方式见 [BlendJob 开发者指南](blendjob.md)。 + +SnapForm 根 package 不再有 `client.py`、`environment.py`、`storage.py` 或 `models.py`。这些文件会制造并不存在的业务边界： + +- Storage Root、`.venv`、Manifest、安装状态与日志路径由 `BlenderJobRuntime` 唯一管理； +- `InstallEnvironment` 由 `runtime.operator_classes()` 提供，Python 与 packages 直接声明在传给 Runtime 的 `environment` 字典中； +- 模型 catalog、磁盘目录、下载和 ONNX Session 都属于 Server，分别位于 `model_catalog.py`、`model_download.py` 与 `model_manager.py`； +- Blender UI 只读取 `model_manager` Resource Snapshot，不导入模型 catalog，也不检查模型文件。 + +`blendjob/client.py` 仍然存在，但它只是 BlendJob 内部的无依赖 HTTP transport，不是 SnapForm 业务抽象，也不向 Operator 暴露 `JobRequest`。 + +SnapForm 的业务 Operator 直接提交参数字典，Job Type 一一对应： + +- `download_model` → `download-model` +- `remove_background` → `remove-background` +- `extract_frames` → `extract-frames` +- `generate_da3_depth` → `generate-da3-depth` +- `generate_moge2_depth` → `generate-moge2-depth` +- `generate_da3_surface` → `generate-da3-surface` +- `generate_moge2_surface` → `generate-moge2-surface` +- `debug_da3` → `debug-da3` +- `debug_moge2` → `debug-moge2` +- `debug_ben2` → `debug-ben2` + +DA3 与 MoGe-2 的生产任务不再合并成带 Backend 参数的通用请求。每个 Operator 的 `request(context)` 直接构造自己的有效参数，Debug 任务也与生产任务完全分开。 + +Server 端没有 `execute_*`、`*_result` 或 CLI worker 中转。全部 `@server.job(...)` 都集中在 `server/__init__.py`：简单任务直接实现，复杂任务调用按领域拆分的 `media.py`、`background.py`、`moge2.py` 或 `da3.py`；共享的纯几何计算位于 `geometry.py`。模型下载由 `ModelManager` 单向调用 `model_download.py`，不存在循环 import。 + +## ONNX-only Environment + +扩展不内置 `uv` 可执行文件或 AI 运行库。`Install Environment` 首先下载固定版本的 uv，校验 SHA-256 后解压到共享 Storage Root，再用它创建环境并安装： + +- ONNX Runtime 1.24.4；Windows 使用 `onnxruntime-directml` +- NumPy 2.4.2、Pillow 12.1.1、PyAV 17.0.0、trimesh 4.11.2 +- FastAPI 0.139.2、Uvicorn 0.51.0、huggingface-hub 1.4.1 + +运行环境不安装 PyTorch、TorchVision、xFormers，也不安装 DA3、MoGe、BEN2 或 utils3d 的 Python 源码包。模型推理和 SnapForm 所需的相机、点云、网格后处理均由 ONNX Runtime、NumPy 与内部实现完成。 + +Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExecutionProvider`。显式指定不可用设备时会安全回退到 CPU。其他平台安装标准 ONNX Runtime，并按可用 provider 选择 CoreML、OpenVINO 或 CPU。 + +`environment.json` 保存规范化后的声明，`manifest.json` 保存该声明与 BlendJob 基础 Server packages 的 Hash。Python、packages、平台 packages 或 BlendJob Environment Schema 改变后需要重新执行 `Install Environment`，不再维护手工 Manifest Revision。 + +## 模型 + +所有正式功能都使用 ONNX： + +- MoGe-2 ViT-S、ViT-B、ViT-L Normal +- DA3 Small、DA3 Base +- BEN2 Base + +首次安装在 `post_install(runtime)` 中自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`。它们和 `Models` panel 的后续下载一样，都是 `download-model` Job，与推理共用队列、HTTP 进度和取消协议。模型按声明的目录名保存在 `~/.snapform/models/`，并保留仓库内部结构。声明了文件清单与校验值的模型先走 Hugging Face `resolve` 直链流式下载，失败后尝试 `https://models.omoolab.xyz//`；两者都失败时当前 Job 直接失败，不记录或猜测后续下载源。建连成功后不限制模型文件读取时长。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 + +`model_manager` Resource 对外统一表示模型系统，持有模型 catalog、磁盘就绪状态和内存 Session。Blender 不拼接模型目录，也不直接检查权重文件。`ModelSessionCache` 保留 BEN2，并在 DA3 与 MoGe-2 之间只保留最近使用的一种 Depth Session。`GET /resources/model_manager` 返回 catalog、下载状态和当前加载状态；`POST /resources/model_manager/clear` 只允许在没有 Active 或 Queued Job 时执行。 + +DA3 ONNX 使用模型的 camera decoder 恢复外参与内参；旧版 `Use Ray Pose` 参数只为任务协议兼容保留，不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics，以及 BEN2 的 Alpha，均在 NumPy 后处理中恢复。 + + +## 安装网络与镜像 + +安装器直接开始工作，不做阻塞式网络预检。packages 由 uv 安装，可用 `UV_INDEX_URL` 显式指定镜像。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 + +uv 默认从 Astral 的 GitHub Release 下载，失败后回退到 `https://mirror.omoolab.xyz/uv//`。两个来源使用相同的连续字节进度并共用固定 SHA-256 校验。下载压缩包只作临时文件，提取完成后删除。模型下载不经过 GitHub。 + +首次安装依次显示 Environment、MoGe-2 下载与 BEN2 下载，每项分别使用 0–100%。Environment 由 BlendJob 依次完成 uv 下载、Python Environment 创建与 packages 安装，uv 不再占用独立一轮进度。两个模型下载由 SnapForm 的 `post_install(runtime)` 通过 `runtime.request()` 顺序发起；切换模型时进度归零，但不显示 `2/2` 等阶段编号。 + +## 生命周期 + +扩展启用后不启动 Server；第一次执行 Job 或点击 `Start` 时才为当前 Blender 启动专属实例。正常关闭或禁用扩展时直接停止自己启动的进程；Server 也监视 Blender 父进程 PID，父进程异常退出后会自行停止。 + +AI、帧提取和 Models panel 下载的 Job 进度保存在 Server 内存中。Handler 只使用 `context.progress(value, message)` 发布进度，终态由 JobServer 内部管理；SnapForm 不再维护独立的 `server/progress.py` 或 Job 状态 JSON。Blender 每 0.25 秒轮询 `GET /jobs/{job_id}`，任务目录不写 `status.json`。Operator 直接提交参数字典；Job 被接受时由 `JobServer` 自动生成 ID 并创建 `jobs//`，提交响应同时返回 `job_id` 与 `directory`。路径通过 `JobContext.directory` 传给 Handler，SnapForm Server 强制将输出写入该目录。 + +`job_runtime.py` 创建唯一的 `runtime = BlenderJobRuntime(...)`，直接接收 `server` 实例并导出 `runtime.JobOperatorBase`。Runtime 独占 Storage Root、Environment 路径、安装进程、Server 控制和 Status Bar。Operator 未覆盖 `request(context)` 时会提交全部 Property；覆盖后返回值纯替换默认参数。可选的 `response(context, result)` 在 Blender 主线程直接完成结果导入。SnapForm 不再提供 Client 请求层、Job Mixin、Status Bar 或 Build 中转 Operator。 + +Environment 安装是 Server 启动前唯一的 Bootstrap 流程，使用 Storage Root 下共享的 `environment-status.json` 和 `environment.log` 回传进度，不创建 Job 目录。Environment 成功后执行普通的 `post_install(runtime)`；默认模型下载此时已经是正式 Server Job，会创建自己的 Job 目录，并使用与其它任务相同的进度和取消协议。 + +Server 使用标准库单 Worker FIFO Queue，全局同时只执行一个 Job，包括模型下载。Server 忙碌时新 Job 进入 `queued`，不再返回 409；排队任务可以在执行前取消。取消运行中任务不会重启 Server；流式模型下载会在分块检查点结束，ONNX Runtime 当前调用则在返回后的最近检查点结束。JobServer 最多保留 100 个已完成 Job 的内存状态，后续提交时清理更早记录。 + +启动与崩溃信息写入 `~/.snapform/server.log`。 + +## 已知限制 + +DA3 ONNX 的注意力是朴素实现（`MatMul → Softmax → MatMul`），峰值显存与 `(帧数 × 每帧 patch 数)²` 成正比。迁移到 ONNX 前的 PyTorch 版使用 `scaled_dot_product_attention`（`XFORMERS_DISABLED=1`），按块计算注意力分数、不落地完整注意力矩阵，显存近似线性，因此视频能处理几十上百帧；ONNX 导出把 SDPA 回退成朴素注意力，多帧在中高分辨率下会显存爆炸（约 12 帧 @ 1008 即需 ~90 GB）。 + +Depth Plane 与 Surface 的视频、序列帧转换会复用同一个 ONNX Session 逐帧独立推理，再把结果合并到现有序列输出，因此峰值显存只取决于单帧分辨率。逐帧结果不保留 DA3 的跨视图尺度与相机位姿一致性，适用于当前逐帧动画结构。Debug 多帧输出仍使用联合推理，以保留多视图关系。 + +联合推理时，`onnx_da3.infer` 会在执行前按 `views × patches` 预估注意力张量，超过 4 GiB 直接报错并提示减少帧数或 Resolution Scale，避免 onnxruntime 报出难以定位的分配失败或 `UnicodeDecodeError`（DirectML 在中文 Windows 下报错信息为 GBK 编码，被 UTF-8 解码掩盖）。 + +待解决方向： + +- 重新导出 DA3 ONNX，使用 ORT 的 `Attention` / `MultiHeadAttention` 贡献算子，或注入内存高效的注意力算子； +- 若未来需要统一坐标系的视频重建，可按小窗口分批推理，再对齐相机位姿与深度尺度。 diff --git a/docs/service.md b/docs/service.md deleted file mode 100644 index 98af95b..0000000 --- a/docs/service.md +++ /dev/null @@ -1,104 +0,0 @@ -# AI Service - -## 运行架构 - -SnapForm 在独立的 Python 3.12 Environment 中运行 HTTP Job Server，Blender 只通过标准库 HTTP Client 提交和轮询任务，不把 AI 运行库载入 Blender Python。Service 只监听 `127.0.0.1`，由 BlendJob `ServiceController` 保存随机端口与 Instance ID；专属本地进程协议不使用额外 Token 鉴权。 - -每个 Blender 实例启动并持有一个 SnapForm 专属服务，不发现或复用其它 Blender、其它插件启动的本地服务。多个 Blender 可以对应多个服务进程，但模型、Environment 与任务产物仍可共享同一个 Storage Root。 - -Environment、模型和任务结果默认保存在 `~/.snapform`： - -```text -~/.snapform/ -├── .venv/ -├── models/ -├── jobs/ -├── manifest.json -├── runtime.json # 持久的模型下载源状态 -├── environment.log -├── environment-status.json -└── service.log -``` - -不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。连接端口与 Instance ID 只保存在启动它们的 Blender 进程内存中。 - -通用 HTTP 层位于 `blendjob` module。`JobServer` 包装 FastAPI 并提供 Job 与 Resource API，`JobClient` 是无第三方依赖的 Blender 侧访问器。SnapForm 的 `ai_service/server.py` 只创建专属 Server、注册任务型 Job，并把 `ModelManager` 注册为 `models` Resource，不直接使用 FastAPI。完整接入方式见 [BlendJob 开发者指南](blendjob.md)。 - -SnapForm 公开的 Request 构造函数与 Job Type 一一对应： - -- `download_model` → `download-model` -- `remove_background` → `remove-background` -- `extract_frames` → `extract-frames` -- `generate_da3_depth` → `generate-da3-depth` -- `generate_moge2_depth` → `generate-moge2-depth` -- `generate_da3_surface` → `generate-da3-surface` -- `generate_moge2_surface` → `generate-moge2-surface` -- `debug_da3` → `debug-da3` -- `debug_moge2` → `debug-moge2` -- `debug_ben2` → `debug-ben2` - -DA3 与 MoGe-2 的生产任务不再合并成带 Backend 参数的通用 Request。每种 Request 只携带自己的有效参数，Debug 任务也与生产任务完全分开。 - -## ONNX-only Environment - -扩展不内置 `uv` 可执行文件或 AI 运行库。`Install Environment` 首先下载固定版本的 uv，校验 SHA-256 后解压到共享 Storage Root，再用它创建环境并安装： - -- ONNX Runtime 1.24.4；Windows 使用 `onnxruntime-directml` -- NumPy 2.4.2、Pillow 12.1.1、PyAV 17.0.0、trimesh 4.11.2 -- FastAPI 0.139.2、Uvicorn 0.51.0、huggingface-hub 1.4.1 - -运行环境不安装 PyTorch、TorchVision、xFormers，也不安装 DA3、MoGe、BEN2 或 utils3d 的 Python 源码包。模型推理和 SnapForm 所需的相机、点云、网格后处理均由 ONNX Runtime、NumPy 与内部实现完成。 - -Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExecutionProvider`。显式指定不可用设备时会安全回退到 CPU。其他平台安装标准 ONNX Runtime，并按可用 provider 选择 CoreML、OpenVINO 或 CPU。 - -`manifest.json` 记录 Environment 版本、ONNX Runtime 版本、runtime variant 与设备列表。Manifest Revision 改变后需要重新执行 `Install Environment`。 - -## 模型 - -所有正式功能都使用 ONNX： - -- MoGe-2 ViT-S、ViT-B、ViT-L Normal -- DA3 Small、DA3 Base -- BEN2 Base - -首次安装由 Bootstrap 流程自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`。Environment 就绪后，`Models` panel 的后续下载提交为 `download-model` Job，与推理共用单任务调度、HTTP 进度和取消协议。模型按声明的目录名保存在 `~/.snapform/models/`，并保留仓库内部结构。声明了文件清单与校验值的模型会直接走 Hugging Face `resolve` 直链流式下载，跳过文件清单与 ETag 预校验；连接源不可达时使用 3 秒 TCP 与 TLS 建连超时，随后回退到 `https://models.omoolab.xyz//`。镜像选择写入持久的 `runtime.json`，后续下载直接走 R2，避免每次 Service 重启后重复等待 Hugging Face。建连成功后不限制模型文件读取时长。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 - -`models` Resource 对外统一表示模型系统，内部区分磁盘模型与内存 Session。`ModelSessionCache` 保留 BEN2，并在 DA3 与 MoGe-2 之间只保留最近使用的一种 Depth Session。`GET /resources/models` 返回当前加载状态；`POST /resources/models/clear` 只允许在没有 Active Job 时执行。 - -DA3 ONNX 使用模型的 camera decoder 恢复外参与内参；旧版 `Use Ray Pose` 参数只为任务协议兼容保留，不参与 ONNX 推理。MoGe-2 的 metric depth、normal、mask、normalized intrinsics，以及 BEN2 的 Alpha，均在 NumPy 后处理中恢复。 - - -## 安装网络与镜像 - -安装器直接开始工作，不做阻塞式网络预检。依赖安装默认依次尝试清华、阿里云 PyPI 镜像并以 pypi.org 兜底，也可用 `UV_INDEX_URL` 显式指定。Windows 系统代理会注入子进程；同时支持 `HTTP_PROXY`、`HTTPS_PROXY` 与 `ALL_PROXY`。 - -uv 默认从 Astral 的 GitHub Release 下载，TCP 与 TLS 建连超过 3 秒时回退到 `https://mirror.omoolab.xyz/uv//`；镜像建连超时为 15 秒，避免较慢的 TLS 握手再次过早失败。两个来源使用相同的下载速度与连续字节进度显示，并共用固定 SHA-256 校验，建连成功后不限制文件读取时长。下载压缩包只作临时文件，提取完成后删除。模型下载不经过 GitHub。 - -首次安装按阶段分别显示 0–100%：uv 下载与解压、ONNX Environment 创建与依赖安装、MoGe-2 下载、BEN2 下载。进入下一阶段时进度条归零；若 uv 已存在则省略该阶段并重新编号。 - -## 生命周期 - -扩展启用后不启动 Service；第一次执行 AI 操作或点击 `Start` 时才为当前 Blender 启动专属实例。正常关闭或禁用扩展时直接停止自己启动的进程；服务也监视 Blender 父进程 PID，父进程异常退出后会自行停止。 - -AI、帧提取和 Models panel 下载的 Job 进度保存在服务内存中。Blender 每 0.25 秒轮询 `GET /jobs/{job_id}`，任务目录不写 `status.json`。`JobRequest` 不包含 ID；Job 被接受时由 `JobServer` 自动生成 ID 并创建 `jobs//`，提交响应同时返回 `job_id` 与 `directory`。路径通过 `JobContext.directory` 传给 handler，SnapForm Service 强制将 Worker 输出写入该目录。 - -SnapForm 的 Blender 侧访问与全部 Request 构造函数集中在 `client.py`。提交线程、轮询、进度、取消和终态清理由 `blendjob.ModalJobOperatorBase` 直接提供；`operators.py` 的 `SnapFormJobOperator` 只保留 SnapForm 输入校验、Request 选取和完成后的 Blender 结果导入。通用运行态是 `ModalJobState`，SnapForm 特有的导入参数单独放在 `SnapFormJobData`。插件专属 Service 只在 `ai_service/server.py` 注册 Job，完全不直接编写 FastAPI Route。 - -只有 Environment 安装及其首次默认模型下载属于 Service 启动前的 Bootstrap 流程，使用 Storage Root 下共享的 `environment-status.json` 和 `environment.log` 回传进度，不创建 Job 目录。 - -Service 全局同时执行一个 Job，包括模型下载。取消任务不会重启 Service；流式模型下载会在分块检查点结束，ONNX Runtime 当前调用则在返回后的最近检查点结束。JobServer 最多保留 100 个已完成 Job 的内存状态，后续提交时清理更早记录。 - -启动与崩溃信息写入 `~/.snapform/service.log`。 - -## 已知限制 - -DA3 ONNX 的注意力是朴素实现（`MatMul → Softmax → MatMul`），峰值显存与 `(帧数 × 每帧 patch 数)²` 成正比。迁移到 ONNX 前的 PyTorch 版使用 `scaled_dot_product_attention`（`XFORMERS_DISABLED=1`），按块计算注意力分数、不落地完整注意力矩阵，显存近似线性，因此视频能处理几十上百帧；ONNX 导出把 SDPA 回退成朴素注意力，多帧在中高分辨率下会显存爆炸（约 12 帧 @ 1008 即需 ~90 GB）。 - -Depth Plane 与 Surface 的视频、序列帧转换会复用同一个 ONNX Session 逐帧独立推理，再把结果合并到现有序列输出，因此峰值显存只取决于单帧分辨率。逐帧结果不保留 DA3 的跨视图尺度与相机位姿一致性，适用于当前逐帧动画结构。Debug 多帧输出仍使用联合推理，以保留多视图关系。 - -联合推理时，`da3_onnx.infer` 会在执行前按 `views × patches` 预估注意力张量，超过 4 GiB 直接报错并提示减少帧数或 Resolution Scale，避免 onnxruntime 报出难以定位的分配失败或 `UnicodeDecodeError`（DirectML 在中文 Windows 下报错信息为 GBK 编码，被 UTF-8 解码掩盖）。 - -待解决方向： - -- 重新导出 DA3 ONNX，使用 ORT 的 `Attention` / `MultiHeadAttention` 贡献算子，或注入内存高效的注意力算子； -- 若未来需要统一坐标系的视频重建，可按小窗口分批推理，再对齐相机位姿与深度尺度。 diff --git a/docs/surface.md b/docs/surface.md index 8b217ad..71dcdeb 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -7,11 +7,11 @@ `Convert to Depth Plane`，当前不处理视频或序列帧。 模型列表包含三个 MoGe-2 Normal 模型和两个 DA3 模型，默认使用 MoGe-2 ViT-S Normal。DA3 的 Resolution Scale 与 Depth Plane 使用相同设计；MoGe-2 改为显示共享的 Resolution Level。 -2. Service 请求 DA3 的 `depth`、`conf` 与 `intrinsics`。仅在开启 `Remove Background` 时运行 BEN2 取得 - 前景 Alpha，并在输出前删除包含背景顶点的三角面；关闭时不加载 BEN2。Depth 不做 smoothing，Service +2. Server 请求 DA3 的 `depth`、`conf` 与 `intrinsics`。仅在开启 `Remove Background` 时运行 BEN2 取得 + 前景 Alpha，并在输出前删除包含背景顶点的三角面；关闭时不加载 BEN2。Depth 不做 smoothing，Server 反投影深度网格，建立源三角面，计算 confidence percentile 与 3×3 相对深度不连续度，并写入 基于原图像素坐标的源 `UVMap`。这些步骤只准备修改器的输入网格，不执行 Relax、拓扑修复、平滑、 - 小岛清理或最终 UV 重投影。运行时显示的六个 `Preparing surface data` 步骤均指这段 Service 端源数据准备。 + 小岛清理或最终 UV 重投影。运行时显示的六个 `Preparing surface data` 步骤均指这段 Server 端源数据准备。 3. Blender 把勾选 `Usage > Modifier` 的 `O DA3 Surface` 节点组直接作为 Geometry Nodes 修改器 添加到 Surface，不另设修改器名称。节点组内没有子节点组，依次完成： - 修改器顶层依次显示 `Resolution`、`Relax Iterations`、`Smooth`；默认展开的 `Culling` Panel @@ -25,7 +25,7 @@ - Smooth 采用 `O Smooth` 的边界检测和模糊法线权重，固定开放边界与薄面，只保留一个 `Smooth` 程度系数；节点组不再额外执行 Flip Faces； - 从修改前的原始表面以 nearest-surface 重新采样 `UVMap`，最终明确写为 Corner 域 `2D Vector`。 - MoGe-2 使用独立的 `O MoGe2 Surface`。Service 由内置 NumPy 网格构建函数 + MoGe-2 使用独立的 `O MoGe2 Surface`。Server 由内置 NumPy 网格构建函数 生成有限深度范围内的紧凑三角网格，并保留模型 mask、逐点 depth 及相邻点相对深度差。Geometry Nodes 按官方 `mask & ~depth_map_edge(depth, rtol=0.04)` 的含义，先由 `Mask` 删除 mask 外的面，再由 `Depth Continuity` 按同一份 3×3 邻域相对深度差删除断层面。DA3 与 MoGe-2 使用相同的属性、 @@ -39,13 +39,13 @@ Group Input，并隐藏没有连线的接口，避免跨越整个节点图的输入连线。 `O DA3 Surface` 与 `O MoGe2 Surface` 的 `Resolution` 默认值均为 `0.1`，所有输入 socket 的 Structure Type 均固定为 `Single`。 -5. Service 准备过程固定显示 6 个编号阶段：验证 DA3 数据、反投影、准备 DA3 与可选 BEN2 数据、构建原始 +5. Server 准备过程固定显示 6 个编号阶段：验证 DA3 数据、反投影、准备 DA3 与可选 BEN2 数据、构建原始 源三角面、源 UV/attributes、写出网格。状态栏显示 `Preparing surface data n/6: ...`。 -6. Service 在 job 根目录写出 `mesh.npz` 与 `texture.png`。NPZ 保存 camera-space `vertices`、 +6. Server 在 job 根目录写出 `mesh.npz` 与 `texture.png`。NPZ 保存 camera-space `vertices`、 `faces`、`uv`、`intrinsics`、`image_size` 与有效深度中位数 `reference_depth`。Blender 读取这些 数据后直接创建 Surface 与原图材质，不需要 Open3D 或 OpenCV 进入 Blender Python 环境。 DA3 Surface 只保留 Geometry Nodes 后处理实际使用的 attributes：点域 `o_depth`、面域 - `o_face_confidence` 与面域 `o_depth_discontinuity`。Remove Background 已在 Service + `o_face_confidence` 与面域 `o_depth_discontinuity`。Remove Background 已在 Server 建面前直接删除背景三角面，不写入任何 background attribute。MoGe-2 的导入网格保留点域 `o_depth`， 以及面域 `o_mask`、`o_depth_discontinuity`，不伪造 confidence、stretch 或 background 属性， 也不写入 custom normal。完整分辨率的 `normal_map` 保存在 NPZ 中并转换为材质贴图。 @@ -56,7 +56,7 @@ 缩放及父级关系产生的世界变换保持一致。更近的顶点沿 Empty 局部 `+Z` 方向凸出，三角面正面朝向 观察相机，启用 Backface Culling 时可从正面看到。 8. 转换对话框不再提供后处理参数，只保留模型与 Resolution Scale 设置、`Create Matching Camera`、 - `Remove Background`、`Unshaded`，以及分割线下方的 `Keep Original`。背景移除默认关闭；开启后 Service + `Remove Background`、`Unshaded`，以及分割线下方的 `Keep Original`。背景移除默认关闭；开启后 Server 使用 BEN2，并在写出 Surface 前直接删除背景三角面，不再由 Geometry Nodes 的 Background 开关处理。 该预删除同时适用于 DA3 与 MoGe-2；因此 MoGe-2 转换后即使关闭 `Mask`，BEN2 已删除的背景也不会恢复。 `Unshaded` 使用与 Depth Plane 相同的 Emission/Transparent 材质路径；普通材质的 @@ -84,7 +84,7 @@ depth、confidence、intrinsics、extrinsics 与 confidence percentile。当前 所有任务使用统一的状态栏进度显示。变化中的阶段说明显示在右侧的进度条内部，取消按钮 位于进度条右侧。状态字段支持阶段、当前操作与预计剩余秒数，界面统一组合为紧凑文案并保证 -进度单调递增。任务进程与 Service 请求在后台启动，按钮点击后立即进入进度状态。Image Depth 的阶段包括 Environment 启动、输入准备、模型加载、深度估算、缓存 +进度单调递增。任务进程与 Server 请求在后台启动，按钮点击后立即进入进度状态。Image Depth 的阶段包括 Environment 启动、输入准备、模型加载、深度估算、缓存 写出和面片构建。进度更新会强制刷新状态栏与窗口，无需移动鼠标即可实时显示。安装中断后残留的 半成品 Environment 会在下次安装前自动清理，不会因为中断而无法再次安装。 @@ -166,9 +166,9 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN `DA3 Nested Giant Large 1.1`。Small 和 Base 没有 Gaussian head，无法由 DA3 生成 `gs_ply`。 -每个 Service Job 被接受时，BlendJob 在 `~/.snapform/jobs/` 下创建一个独立文件夹。文件夹名是 32 位 opaque Job ID，不包含任务类型或对象名称。执行逻辑从 `JobContext.directory` 获取路径，输出产物直接写入该目录。 +每个 Server Job 被接受时，BlendJob 在 Runtime Storage Root 的 `jobs/` 下创建一个独立文件夹。文件夹名是 32 位 opaque Job ID，不包含任务类型或对象名称。执行逻辑从 `JobContext.directory` 获取路径，输出产物直接写入该目录。 -Blender 为 Packed Image 或序列帧生成的输入副本先保存在系统临时目录，Service 读取完成后由 Blender 清理，不归档到 Job 目录。普通磁盘文件和视频仍直接读取原文件。Service Job 日志统一写入 `service.log`，进度只保存在 Service 内存中。Environment Bootstrap 使用 Storage Root 下的 `environment.log` 与 `environment-status.json`，不创建 Job 目录。 +Blender 为 Packed Image 或序列帧生成的输入副本先保存在系统临时目录，Server 读取完成后由 Blender 清理，不归档到 Job 目录。普通磁盘文件和视频仍直接读取原文件。Server 日志统一写入 `server.log`，进度只保存在 Server 内存中。Environment Bootstrap 使用 Storage Root 下的 `environment.log` 与 `environment-status.json`，不创建 Job 目录。 与几何输出相关的 Job Type 为 `generate-da3-depth`、`generate-moge2-depth`、`generate-da3-surface`、`generate-moge2-surface`、`debug-da3` 与 `debug-moge2`。每种 Job 只接收对应模型与输出所需的参数。 diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 2641f33..b2b6b65 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -1,29 +1,21 @@ import bpy -from . import client from .blender import depth_plane as image_depth from .blender import surface as image_mesh from .blender.depth_plane import ( - ApplyRemovedBackground, - BuildDepthPlane, - BuildOriginalDepthPlane, - BuildMoge2DepthPlane, - ConvertImageToDepthPlane, + ConvertToDepthPlane, RemoveImageBackground, SnapFormImageMenu, draw_image_context_menu, ) -from .blender.surface import BuildImageMesh, ConvertImageToMesh +from .blender.surface import ConvertToSurface from .operators import ( - CancelJob, - ControlService, JOB_OPERATOR_CLASSES, - cancel_active_job, - close_active_job, ) +from .job_runtime import runtime from .panel import ( Ben2RemovalDebugPanel, - AiServicePanel, + ServerPanel, Da3OutputDebugPanel, DebugPanel, ModelsPanel, @@ -31,25 +23,17 @@ from .panel import ( ) from .preferences import SnapFormPreferences from .properties import SnapFormSettings -from .status_bar import draw_status_bar CLASSES = ( SnapFormPreferences, SnapFormSettings, *JOB_OPERATOR_CLASSES, - CancelJob, - ControlService, SnapFormImageMenu, - ApplyRemovedBackground, RemoveImageBackground, - BuildDepthPlane, - BuildOriginalDepthPlane, - BuildMoge2DepthPlane, - ConvertImageToDepthPlane, - BuildImageMesh, - ConvertImageToMesh, - AiServicePanel, + ConvertToDepthPlane, + ConvertToSurface, + ServerPanel, ModelsPanel, DebugPanel, Moge2ApiDebugPanel, @@ -64,20 +48,15 @@ def register(): bpy.types.Scene.snapform_settings = bpy.props.PointerProperty( type=SnapFormSettings, ) - bpy.types.STATUSBAR_HT_header.append(draw_status_bar) bpy.types.VIEW3D_MT_object_context_menu.append(draw_image_context_menu) bpy.types.OUTLINER_MT_object.append(draw_image_context_menu) - client.schedule_service_start() + runtime.register() def unregister(): - cancel_active_job() - close_active_job() - client.unschedule_service() - client.detach_service_client() + runtime.unregister() bpy.types.VIEW3D_MT_object_context_menu.remove(draw_image_context_menu) bpy.types.OUTLINER_MT_object.remove(draw_image_context_menu) - bpy.types.STATUSBAR_HT_header.remove(draw_status_bar) if hasattr(bpy.types.Scene, "snapform_settings"): del bpy.types.Scene.snapform_settings for class_type in reversed(CLASSES): diff --git a/src/snapform/ai_service/__init__.py b/src/snapform/ai_service/__init__.py deleted file mode 100644 index 540ce8f..0000000 --- a/src/snapform/ai_service/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""SnapForm AI Service implementation without Blender dependencies.""" diff --git a/src/snapform/ai_service/jobs.py b/src/snapform/ai_service/jobs.py deleted file mode 100644 index b31f22c..0000000 --- a/src/snapform/ai_service/jobs.py +++ /dev/null @@ -1,2277 +0,0 @@ -import argparse -import hashlib -import http.client -import json -import os -import shutil -import socket -import sys -import tempfile -import threading -import time -from pathlib import Path -from urllib.parse import quote -from urllib.request import ( - HTTPSHandler, - ProxyHandler, - Request, - build_opener, - urlopen, -) - - -SCRIPT_DIRECTORY = Path(__file__).resolve().parent -if str(SCRIPT_DIRECTORY) not in sys.path: - sys.path.insert(0, str(SCRIPT_DIRECTORY)) - - -IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} -VIDEO_SUFFIXES = { - ".mp4", - ".mov", - ".avi", - ".mkv", - ".webm", - ".m4v", - ".mts", - ".mpg", - ".mpeg", -} -SCENE_FILENAME = "scene.glb" -EXPORT_FORMATS = ( - "glb", - "gs_ply", - "depth_vis", - "mini_npz", - "mesh", - "mesh_data", -) -HF_ENDPOINT = "https://huggingface.co" -R2_ENDPOINT = "https://models.omoolab.xyz" -# 连接（TCP + TLS 握手）超时，源不可达时快速失败切换；下载阶段不设读超时。 -CONNECT_TIMEOUT = 3.0 -DOWNLOAD_CHUNK_SIZE = 256 * 1024 - - -class _ShortConnectHTTPSConnection(http.client.HTTPSConnection): - """HTTPS 连接：TCP + TLS 握手用短超时，读操作不设超时。""" - - def connect(self): - try: - self.timeout = CONNECT_TIMEOUT - super().connect() - finally: - self.timeout = None - # 连接建立后读操作不设超时，下载大文件不会被读超时中断。 - self.sock.settimeout(None) - - -class _ShortConnectHTTPSHandler(HTTPSHandler): - def https_open(self, req): - return self.do_open(_ShortConnectHTTPSConnection, req) - - -_download_opener = build_opener(ProxyHandler(), _ShortConnectHTTPSHandler()) - - -class TaskCancelled(RuntimeError): - pass - - -def check_cancelled(cancel_check): - if cancel_check is not None and cancel_check(): - raise TaskCancelled("Task cancelled") - - -def write_status( - target, - progress, - message, - error=None, - state=None, - stage=None, - stages=None, - stage_label=None, - eta_seconds=None, -): - payload = {"progress": progress, "message": message} - if error: - payload["error"] = error - if stage is not None: - payload["stage"] = stage - payload["stages"] = stages - payload["stage_label"] = stage_label - if eta_seconds is not None: - payload["eta_seconds"] = max(int(eta_seconds), 0) - if state is None: - if error: - state = "failed" - elif progress >= 1.0: - state = "succeeded" - else: - state = "running" - payload["state"] = state - if callable(target): - target( - progress, - message, - error=error, - state=state, - stage=stage, - stages=stages, - stage_label=stage_label, - eta_seconds=eta_seconds, - ) - return - Path(target).write_text( - json.dumps(payload, ensure_ascii=False), - encoding="utf-8", - ) - - -def collect_images(directory): - directory = Path(directory) - return sorted( - path - for path in directory.iterdir() - if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES - ) - - -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, - av, - start_frame=1, - cancel_check=None, -): - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - try: - container = av.open(str(video_path)) - except Exception as error: - raise RuntimeError(f"Unable to open video: {video_path}") from error - image_paths = [] - skipped = 0 - try: - stream = container.streams.video[0] - for frame in container.decode(stream): - check_cancelled(cancel_check) - if skipped < start_frame - 1: - skipped += 1 - continue - image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" - frame.to_image().save(image_path, format="JPEG") - image_paths.append(image_path) - if len(image_paths) >= maximum_frames: - break - finally: - container.close() - - 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, - 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 - if suffix not in VIDEO_SUFFIXES: - supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) - raise RuntimeError(f"Unsupported input type; expected one of: {supported}") - - import av - - temporary_directory = Path( - tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root) - ) - try: - images = extract_video_frames( - input_path, - temporary_directory, - video_frames, - av, - start_frame=start_frame, - cancel_check=cancel_check, - ) - except Exception: - shutil.rmtree(temporary_directory, ignore_errors=True) - raise - 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" - destination = ( - Path(output_directory) / f"frame.{index + 1:04d}{suffix}" - ) - if Path(image_path).resolve() == destination.resolve(): - continue - shutil.copy2(image_path, destination) - - -def unproject_depth_frame(depth, intrinsics, extrinsics): - import numpy as np - - height, width = depth.shape - x, y = np.meshgrid( - np.arange(width, dtype=np.float32), - np.arange(height, dtype=np.float32), - ) - pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) - rays = (np.linalg.inv(intrinsics) @ pixels.T).T - camera_points = rays * depth.reshape(-1, 1) - camera_points = np.column_stack( - (camera_points, np.ones(camera_points.shape[0], dtype=np.float32)) - ) - world_to_camera = np.eye(4, dtype=np.float32) - world_to_camera[: extrinsics.shape[0], :4] = extrinsics - world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T - return world_points[:, :3].reshape(height, width, 3) - - -def unproject_camera_depth(depth, intrinsics): - import numpy as np - - height, width = depth.shape - x, y = np.meshgrid( - np.arange(width, dtype=np.float32), - np.arange(height, dtype=np.float32), - ) - pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) - rays = (np.linalg.inv(intrinsics) @ pixels.T).T - return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) - - -def compact_mesh(vertices, faces): - import numpy as np - - vertices = np.asarray(vertices, dtype=np.float32) - faces = np.asarray(faces, dtype=np.int32) - used_vertices = np.unique(faces) - index_map = np.full(len(vertices), -1, dtype=np.int32) - index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32) - vertices = vertices[used_vertices] - faces = index_map[faces] - return vertices, faces - - -def project_camera_uv(vertices, intrinsics, image_size): - import numpy as np - - width, height = image_size - z = vertices[:, 2] - if np.any(~np.isfinite(z)) or np.any(z <= 0.0): - raise ValueError("Simplified mesh contains invalid camera depth") - u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2] - v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2] - denominator_x = max(width - 1, 1) - denominator_y = max(height - 1, 1) - return np.column_stack( - ( - np.clip(u / denominator_x, 0.0, 1.0), - np.clip(1.0 - v / denominator_y, 0.0, 1.0), - ) - ).astype(np.float32) - - -def project_camera_pixels(vertices, intrinsics): - import numpy as np - - z = vertices[:, 2] - if np.any(~np.isfinite(z)) or np.any(z <= 0.0): - raise ValueError("Mesh contains invalid camera depth") - return np.column_stack( - ( - intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2], - intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2], - ) - ) - - -def sample_image_nearest(values, pixels): - import numpy as np - - height, width = values.shape - x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1) - y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1) - return np.asarray(values)[y, x] - - -def confidence_percentile_map(depth, confidence): - import numpy as np - - valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) - result = np.zeros(depth.shape, dtype=np.float32) - if not valid.any(): - return result - sorted_values = np.sort(np.asarray(confidence)[valid]) - ranks = np.searchsorted( - sorted_values, - np.asarray(confidence)[valid], - side="right", - ) - result[valid] = ranks.astype(np.float32) / len(sorted_values) - return result - - -def depth_discontinuity_ratio(depth): - import numpy as np - - padded = np.pad(depth, 1, mode="edge") - windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) - difference = np.nanmax(windows, axis=(-2, -1)) - difference -= np.nanmin(windows, axis=(-2, -1)) - return difference / np.maximum(depth, 1e-6) - - -def triangle_edge_stretch_values(points, triangles): - import numpy as np - - flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3) - triangle_points = flat_points[triangles] - triangle_depth = triangle_points[..., 2] - safe_depth = np.where( - np.abs(triangle_depth) > 1e-8, - triangle_depth, - np.nan, - ) - triangle_rays = triangle_points / safe_depth[..., None] - next_points = np.roll(triangle_points, -1, axis=1) - next_rays = np.roll(triangle_rays, -1, axis=1) - next_depth = np.roll(triangle_depth, -1, axis=1) - local_depth = (triangle_depth + next_depth) * 0.5 - flat_edges = ( - triangle_rays * local_depth[..., None] - - next_rays * local_depth[..., None] - ) - flat_lengths = np.linalg.norm(flat_edges, axis=-1) - actual_lengths = np.linalg.norm( - triangle_points - next_points, - axis=-1, - ) - stretch = actual_lengths / np.maximum(flat_lengths, 1e-8) - stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0 - return stretch - - -def image_mesh_attributes( - vertices, - faces, - intrinsics, - confidence_percentiles, - face_depth_discontinuity, -): - import numpy as np - - point_percentile = sample_image_nearest( - confidence_percentiles, - project_camera_pixels(vertices, intrinsics), - ).astype(np.float32) - face_confidence = point_percentile[faces].min(axis=1) - return { - "point_depth": vertices[:, 2].astype(np.float32), - "face_confidence_percentile": face_confidence.astype(np.float32), - "face_depth_discontinuity": np.asarray( - face_depth_discontinuity, - dtype=np.float32, - ), - } - - -def export_prediction_attributes(prediction, output_directory): - import numpy as np - - values = {} - for name in ("depth", "conf", "intrinsics", "extrinsics"): - value = getattr(prediction, name, None) - if value is not None: - values[name] = np.asarray(value) - if "depth" in values and "conf" in values: - values["confidence_percentile"] = np.stack( - [ - confidence_percentile_map(depth, confidence) - for depth, confidence in zip( - values["depth"], - values["conf"], - ) - ] - ) - output_path = Path(output_directory) / "prediction.attributes.npz" - np.savez(output_path, **values) - return output_path - - -def export_da3_mini_npz(prediction, output_directory): - import numpy as np - - values = {"depth": np.round(prediction.depth, 8)} - for name in ("conf", "extrinsics", "intrinsics"): - value = getattr(prediction, name, None) - if value is not None: - values[name] = np.round(value, 2) if name == "conf" else value - output_path = Path(output_directory) / "results.npz" - np.savez_compressed(output_path, **values) - return output_path - - -def export_da3_depth_preview(prediction, output_directory): - import numpy as np - from PIL import Image - - target = Path(output_directory) / "depth_vis" - target.mkdir(parents=True, exist_ok=True) - for index, depth in enumerate(prediction.depth): - valid = np.isfinite(depth) & (depth > 0.0) - normalized = np.zeros(depth.shape, dtype=np.float32) - if valid.any(): - low, high = np.percentile(depth[valid], (2.0, 98.0)) - normalized[valid] = np.clip( - (depth[valid] - low) / max(float(high - low), 1e-6), - 0.0, - 1.0, - ) - color = np.stack( - ( - normalized, - 1.0 - np.abs(normalized * 2.0 - 1.0), - 1.0 - normalized, - ), - axis=-1, - ) - color = np.asarray(color * 255.0, dtype=np.uint8) - source = np.asarray(prediction.processed_images[index], dtype=np.uint8) - combined = np.concatenate((source, color), axis=1) - Image.fromarray(combined, mode="RGB").save( - target / f"{index:04d}.jpg", - quality=95, - ) - - -def export_da3_point_cloud( - prediction, - output_directory, - confidence_percentile=40.0, - num_max_points=1_000_000, -): - import numpy as np - import trimesh - - depth = np.asarray(prediction.depth, dtype=np.float32) - confidence = np.asarray(prediction.conf, dtype=np.float32) - threshold = confidence_percentile_threshold( - depth, - confidence, - confidence_percentile, - ) - point_sets = [] - color_sets = [] - for index in range(depth.shape[0]): - points = unproject_depth_frame( - depth[index], - prediction.intrinsics[index], - prediction.extrinsics[index], - ) - valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0) - if threshold is not None: - valid &= confidence[index] >= threshold - point_sets.append(points[valid]) - color_sets.append(prediction.processed_images[index][valid]) - points = np.concatenate(point_sets, axis=0) - colors = np.concatenate(color_sets, axis=0) - maximum = max(int(num_max_points), 1) - if len(points) > maximum: - indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64) - points = points[indices] - colors = colors[indices] - scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors)) - scene.export(Path(output_directory) / SCENE_FILENAME) - - -def remove_background_faces(faces, foreground_mask): - """Return only triangles whose vertices all belong to the foreground.""" - foreground_points = foreground_mask.reshape(-1) - return faces[foreground_points[faces].all(axis=1)] - - -def export_image_mesh_data( - prediction, - image_path, - output_directory, - foreground_path=None, - progress_callback=None, -): - import numpy as np - from PIL import Image - - total_steps = 6 - - def progress(step, label): - if progress_callback is not None: - progress_callback(step, total_steps, label) - - progress(1, "Validating DA3 data") - required = ("depth", "conf", "intrinsics") - missing = [ - name - for name in required - if getattr(prediction, name, None) is None - ] - if missing: - raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") - depth_frames = np.asarray(prediction.depth, dtype=np.float32) - confidence_frames = np.asarray(prediction.conf, dtype=np.float32) - if depth_frames.shape[0] != 1: - raise ValueError("Convert to Surface currently supports one image") - raw_depth = depth_frames[0] - confidence = confidence_frames[0] - intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32) - progress(2, "Unprojecting the depth grid") - points = unproject_camera_depth(raw_depth, intrinsics) - progress(3, "Preparing confidence and background data") - foreground_mask = np.ones(raw_depth.shape, dtype=bool) - if foreground_path is not None: - foreground_mask = foreground_mesh_mask( - foreground_path, - raw_depth.shape, - ) - confidence_percentiles = confidence_percentile_map( - raw_depth, - confidence, - ) - progress(4, "Building source triangles") - faces = mesh_triangle_faces( - points, - raw_depth, - confidence, - None, - False, - 0.0, - ) - if foreground_path is not None: - faces = remove_background_faces(faces, foreground_mask) - if not len(faces): - raise ValueError( - "DA3 output does not contain any valid surface faces" - ) - progress(5, "Preparing source UV and attributes") - depth_discontinuity = depth_discontinuity_ratio(raw_depth) - face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) - vertices, faces = compact_mesh( - points.reshape(-1, 3), - faces, - ) - height, width = raw_depth.shape - uv = project_camera_uv(vertices, intrinsics, (width, height)) - attributes = image_mesh_attributes( - vertices, - faces, - intrinsics, - confidence_percentiles, - face_depth_discontinuity, - ) - reference_depth = float(np.median(vertices[:, 2])) - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - progress(6, "Writing source surface data") - result_path = output_directory / "mesh.npz" - np.savez_compressed( - result_path, - vertices=vertices, - faces=faces, - uv=uv, - intrinsics=intrinsics, - image_size=np.asarray((width, height), dtype=np.int32), - reference_depth=np.asarray(reference_depth, dtype=np.float32), - **attributes, - ) - texture_path = output_directory / "texture.png" - texture_source = foreground_path or image_path - with Image.open(texture_source) as image: - image.convert("RGBA").save(texture_path, format="PNG", compress_level=1) - return result_path, texture_path - - -def depth_edge_mask(depth, threshold=12.0 / 255.0): - import numpy as np - - values = np.asarray(depth, dtype=np.float32) - edges = np.zeros(values.shape, dtype=bool) - for axis in (0, 1): - first = np.take(values, range(values.shape[axis] - 1), axis=axis) - second = np.take(values, range(1, values.shape[axis]), axis=axis) - valid = ( - np.isfinite(first) - & np.isfinite(second) - & (first > 0.0) - & (second > 0.0) - ) - scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) - discontinuity = valid & (np.abs(first - second) / scale > threshold) - lower = [slice(None), slice(None)] - upper = [slice(None), slice(None)] - lower[axis] = slice(0, -1) - upper[axis] = slice(1, None) - edges[tuple(lower)] |= discontinuity - edges[tuple(upper)] |= discontinuity - return edges - - -def edge_culling_stretch_limit(culling): - factor = min(max(float(culling) / 100.0, 0.0001), 1.0) - return 1.0 / factor - - -def foreground_mesh_mask(foreground_path, depth_shape): - import numpy as np - from PIL import Image, ImageFilter - - height, width = depth_shape - with Image.open(foreground_path) as source: - alpha = source.convert("RGBA").getchannel("A") - alpha = alpha.resize( - (width, height), - resample=Image.Resampling.BILINEAR, - ) - alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75)) - return np.asarray(alpha) >= 128 - - -def confidence_percentile_threshold(depth, confidence, percentile): - import numpy as np - - if float(percentile) <= 0.0: - return None - valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) - if not valid.any(): - raise ValueError("DA3 returned no valid depth confidence values") - return float(np.percentile(confidence[valid], percentile)) - - -def mesh_triangle_faces( - points, - depth, - confidence, - confidence_threshold, - filter_edges, - max_edge_length, - edge_threshold=12.0 / 255.0, - edge_stretch_limit=None, - vertex_mask=None, -): - import numpy as np - - height, width = depth.shape - indices = np.arange(height * width, dtype=np.int32).reshape(height, width) - quads = np.stack( - ( - indices[:-1, :-1], - indices[1:, :-1], - indices[1:, 1:], - indices[:-1, 1:], - ), - axis=-1, - ).reshape(-1, 4) - valid = ( - np.isfinite(points).all(axis=-1) - & np.isfinite(depth) - & (depth > 0.0) - ) - if confidence_threshold is not None: - valid &= np.isfinite(confidence) & ( - confidence >= confidence_threshold - ) - if filter_edges: - valid &= ~depth_edge_mask(depth, edge_threshold) - if vertex_mask is not None: - if np.shape(vertex_mask) != depth.shape: - raise ValueError("Foreground mask and depth dimensions do not match") - valid &= np.asarray(vertex_mask, dtype=bool) - triangles = np.concatenate( - (quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), - axis=0, - ) - face_valid = valid.reshape(-1)[triangles].all(axis=1) - if edge_stretch_limit is not None: - face_valid &= triangle_edge_stretch_mask( - points, - triangles, - edge_stretch_limit, - ) - if max_edge_length > 0.0: - flat_points = points.reshape(-1, 3) - triangle_points = flat_points[triangles] - next_points = np.roll(triangle_points, -1, axis=1) - edge_lengths = np.linalg.norm( - triangle_points - next_points, - axis=-1, - ) - face_valid &= (edge_lengths <= max_edge_length).all(axis=1) - triangles = triangles[face_valid] - if not len(triangles): - return np.empty((0, 3), dtype=np.int32) - return triangles - - -def triangle_edge_stretch_mask(points, triangles, stretch_limit): - stretch = triangle_edge_stretch_values(points, triangles) - tolerance = float(stretch_limit) * (1.0 + 1e-5) - return (stretch <= tolerance).all(axis=1) - - -def export_textured_mesh_glb( - prediction, - image_paths, - output_directory, - confidence_percentile, - filter_edges, - max_edge_length, -): - import numpy as np - import trimesh - from PIL import Image - - required = ("depth", "conf", "intrinsics", "extrinsics") - missing = [name for name in required if getattr(prediction, name) is None] - if missing: - raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") - depth = np.asarray(prediction.depth, dtype=np.float32) - confidence = np.asarray(prediction.conf, dtype=np.float32) - threshold = confidence_percentile_threshold( - depth, - confidence, - confidence_percentile, - ) - meshes = [] - referenced_points = [] - for index in range(depth.shape[0]): - points = unproject_depth_frame( - depth[index], - prediction.intrinsics[index], - prediction.extrinsics[index], - ) - faces = mesh_triangle_faces( - points, - depth[index], - confidence[index], - threshold, - filter_edges, - max_edge_length, - ) - if not len(faces): - continue - height, width = depth[index].shape - u = np.linspace(0.0, 1.0, width, dtype=np.float32) - v = np.linspace(1.0, 0.0, height, dtype=np.float32) - uu, vv = np.meshgrid(u, v) - uv = np.stack((uu, vv), axis=-1).reshape(-1, 2) - with Image.open(image_paths[index]) as source_image: - texture_image = source_image.convert("RGB").copy() - visual = trimesh.visual.texture.TextureVisuals( - uv=uv, - image=texture_image, - ) - mesh = trimesh.Trimesh( - vertices=points.reshape(-1, 3), - faces=faces, - visual=visual, - process=False, - ) - mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1) - mesh.vertex_attributes["o_confidence"] = confidence[ - index - ].reshape(-1) - mesh.vertex_attributes["o_confidence_percentile"] = ( - confidence_percentile_map( - depth[index], - confidence[index], - ).reshape(-1) - ) - mesh.remove_unreferenced_vertices() - meshes.append(mesh) - referenced_points.append(mesh.vertices) - if not meshes: - raise ValueError( - "Surface filters removed every face; lower Confidence Threshold or disable filters" - ) - - world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64) - if world_to_camera.shape == (3, 4): - homogeneous = np.eye(4, dtype=np.float64) - homogeneous[:3, :4] = world_to_camera - world_to_camera = homogeneous - cv_to_gltf = np.eye(4, dtype=np.float64) - cv_to_gltf[1, 1] = -1.0 - cv_to_gltf[2, 2] = -1.0 - alignment = cv_to_gltf @ world_to_camera - aligned_points = trimesh.transform_points( - np.concatenate(referenced_points, axis=0), - alignment, - ) - center = np.median(aligned_points, axis=0) - alignment[:3, 3] -= center - - scene = trimesh.Scene() - for index, mesh in enumerate(meshes): - mesh.apply_transform(alignment) - scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}") - output_directory = Path(output_directory) - output_directory.mkdir(parents=True, exist_ok=True) - output_path = output_directory / "mesh.glb" - scene.export(output_path) - return output_path - - -class DownloadProgressReporter: - def __init__( - self, - status_path, - repository, - clock=time.monotonic, - start=0.05, - span=0.90, - stage=None, - stages=None, - stage_label=None, - label=None, - ): - self.status_path = status_path - self.repository = repository - self.clock = clock - self.lock = threading.Lock() - self.start = start - self.span = span - self.stage = stage - self.stages = stages - self.stage_label = stage_label - self.label = label - self.progress = start - self.last_write_time = 0.0 - self.started_at = self.clock() - - def report(self, current, total, rate=None): - if not total or current <= 0: - return - - factor = min(max(current / total, 0.0), 1.0) - progress = self.start + factor * self.span - now = self.clock() - with self.lock: - if progress <= self.progress: - return - if now - self.last_write_time < 0.20 and factor < 1.0: - return - self.progress = progress - self.last_write_time = now - effective_rate = rate - elapsed = now - self.started_at - if not effective_rate and elapsed > 0: - effective_rate = current / elapsed - speed = format_rate(effective_rate) if effective_rate else "" - prefix = self.label or self.repository - message = f"Downloading {prefix}" - if speed: - message = f"{message} {speed}" - write_status( - self.status_path, - progress, - message, - stage=self.stage, - stages=self.stages, - stage_label=self.stage_label, - ) - - -def create_download_progress_class( - status_path, - repository, - tqdm_type=None, - clock=time.monotonic, - start=0.05, - span=0.90, - stage=None, - stages=None, - stage_label=None, - label=None, -): - if tqdm_type is None: - from tqdm.auto import tqdm as tqdm_type - - reporter = DownloadProgressReporter( - status_path, - repository, - clock, - start, - span, - stage, - stages, - stage_label, - label, - ) - - class DownloadProgress(tqdm_type): - def __init__(self, *args, **kwargs): - description = str(kwargs.get("desc", "")) - # huggingface_hub 1.x 的字节进度条描述是 "Downloading (incomplete total...)"， - # 用前缀匹配而非精确匹配，否则真实下载进度会被忽略。 - self.snapform_reports_progress = ( - description.startswith("Downloading") - or description.startswith("Fetching ") - ) - super().__init__(*args, **kwargs) - - def update(self, amount=1): - result = super().update(amount) - if self.snapform_reports_progress: - reporter.report(self.n, self.total, self._download_rate()) - return result - - def refresh(self, *args, **kwargs): - result = super().refresh(*args, **kwargs) - if self.snapform_reports_progress: - reporter.report(self.n, self.total, self._download_rate()) - return result - - def _download_rate(self): - try: - return getattr(self, "format_dict", {}).get("rate") - except Exception: - return None - - return DownloadProgress - - -def progress_range(args, staged): - start = getattr(args, "progress_start", None) - end = getattr(args, "progress_end", None) - if start is not None and end is not None: - return float(start), float(end) - float(start) - if staged: - return 0.0, 1.0 - return 0.05, 0.90 - - -def _runtime_state_path(model_dir): - """Return the storage-level download routing state path.""" - return Path(model_dir).parent.parent / "runtime.json" - - -def _models_from_r2(model_dir): - try: - payload = json.loads( - _runtime_state_path(model_dir).read_text(encoding="utf-8") - ) - except (OSError, ValueError, json.JSONDecodeError): - return False - return bool(payload.get("models_from_r2")) - - -def _mark_models_from_r2(model_dir): - """记录当前运行周期的模型镜像选择。""" - path = _runtime_state_path(model_dir) - try: - payload = json.loads(path.read_text(encoding="utf-8")) - except (OSError, ValueError, json.JSONDecodeError): - payload = {} - payload["models_from_r2"] = True - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f"{path.name}.tmp") - temporary.write_text(json.dumps(payload), encoding="utf-8") - temporary.replace(path) - - -def download_model(args): - label = getattr(args, "label", None) or args.repository - model_dir = Path(args.model_dir) - model_dir.mkdir(parents=True, exist_ok=True) - staged = getattr(args, "stage", None) is not None - start, span = progress_range(args, staged) - stage_kwargs = { - "stage": getattr(args, "stage", None), - "stages": getattr(args, "stages", None), - "stage_label": getattr(args, "stage_label", None), - } - r2_directory = getattr(args, "r2_directory", "") - r2_files = getattr(args, "r2_file", []) - if r2_directory and r2_files: - _download_declared_or_fail( - args, - model_dir, - label, - start, - span, - stage_kwargs, - ) - return - _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs) - - -def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs): - """Stream declared files from Hugging Face, then the R2 mirror, remembering fallback.""" - if _models_from_r2(model_dir): - sources = ( - ( - R2_ENDPOINT, - download_from_r2, - f"Downloading {label} (mirror)", - ), - ) - else: - sources = ( - (HF_ENDPOINT, download_from_hf, f"Downloading {label}"), - ( - R2_ENDPOINT, - download_from_r2, - f"Downloading {label} (mirror)", - ), - ) - failures = [] - for position, (source, downloader, message) in enumerate(sources): - write_status(args.status, start, message, **stage_kwargs) - try: - downloader( - args, - model_dir, - start=start, - span=span, - label=label, - stage_kwargs=stage_kwargs, - ) - _verify_model_complete( - model_dir, - getattr(args, "required_pattern", None), - ) - except Exception as error: - cancel_check = getattr(args, "cancel_check", None) - if cancel_check is not None: - cancel_check() - failures.append((source, error)) - continue - # 非首个源成功意味着 hf 已失败，记住本次 service 周期后续直接走 r2。 - if position > 0: - _mark_models_from_r2(model_dir) - write_status( - args.status, - start + span, - f"{label} ready", - **stage_kwargs, - ) - return - detail = "; ".join( - f"{source}: {error}" for source, error in failures - ) - raise RuntimeError(f"All download sources failed: {detail}") - - -def _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs): - """Fallback download for repositories without declared files, via huggingface_hub.""" - failures = [] - try: - from huggingface_hub import snapshot_download - except Exception as error: - snapshot_download = None - failures.append((HF_ENDPOINT, error)) - progress_class = None - if snapshot_download: - progress_class = create_download_progress_class( - args.status, - args.repository, - start=start, - span=span, - label=label, - **stage_kwargs, - ) - for endpoint in download_endpoints() if snapshot_download else (): - os.environ["HF_ENDPOINT"] = endpoint - # snapshot_download 先拉取文件清单并校验 ETag，还没有字节下载， - # 用 Checking 提示当前阶段，避免进度条看起来卡住。 - message = f"Checking {label} files" - 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, - getattr(args, "required_pattern", None), - ) - except Exception as error: - cancel_check = getattr(args, "cancel_check", None) - if cancel_check is not None: - cancel_check() - failures.append((endpoint, error)) - continue - write_status( - args.status, - start + span, - f"{label} ready", - **stage_kwargs, - ) - return - detail = "; ".join( - f"{source}: {error}" for source, error in failures - ) - raise RuntimeError(f"All download sources failed: {detail}") - - -def download_endpoints(): - """Return the configured Hugging Face endpoint or official source.""" - endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT - return [endpoint.rstrip("/")] - - -def _verify_model_complete(model_dir, required_patterns=None): - """Refuse to report success when the required model files are missing.""" - model_dir = Path(model_dir) - required_patterns = required_patterns or ( - "config.json", - "*.safetensors", - ) - missing = [ - pattern - for pattern in required_patterns - if not any(model_dir.glob(pattern)) - ] - if missing: - raise RuntimeError( - "Model download reported success but files are missing: " - + ", ".join(missing) - ) - - -def _download_file_from_url( - url, - destination, - expected_size, - expected_checksum, - progress=None, - retries=2, - cancel_check=None, -): - """Download and verify one model file atomically from a direct URL.""" - request = Request(url, headers={"User-Agent": "SnapForm"}) - opener = _download_opener.open if request.type == "https" else urlopen - last_error = None - for _attempt in range(retries): - temporary = destination.with_name(f"{destination.name}.part") - try: - checksum = hashlib.sha256() - downloaded = 0 - with opener(request) as response: - with open(temporary, "wb") as output: - while True: - if cancel_check is not None: - cancel_check() - chunk = response.read(DOWNLOAD_CHUNK_SIZE) - if not chunk: - break - output.write(chunk) - checksum.update(chunk) - downloaded += len(chunk) - if progress is not None: - progress(len(chunk)) - if downloaded != expected_size: - raise RuntimeError( - f"Size mismatch for {destination.name}: " - f"{downloaded} != {expected_size}" - ) - if checksum.hexdigest().lower() != expected_checksum.lower(): - raise RuntimeError(f"SHA-256 mismatch for {destination.name}") - temporary.replace(destination) - return - except Exception as error: - if cancel_check is not None: - cancel_check() - last_error = error - try: - temporary.unlink() - except OSError: - pass - raise last_error - - -def r2_file_url(directory, filename): - return f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}" - - -def hf_file_url(repository, filename): - """Hugging Face resolve URL for a direct, metadata-free download.""" - endpoint = (os.environ.get("HF_ENDPOINT") or HF_ENDPOINT).rstrip("/") - return f"{endpoint}/{repository}/resolve/main/{quote(filename)}" - - -def _parse_declared_files(args): - selected = [] - for value in args.r2_file: - filename, size, checksum = value.split("|", 2) - selected.append((filename, int(size), checksum)) - return selected - - -def _download_declared_files( - args, - model_dir, - url_builder, - start, - span, - label, - stage_kwargs, - retries=2, -): - """Download each declared file through url_builder with aggregated progress.""" - repository = args.repository - selected = _parse_declared_files(args) - reporter = DownloadProgressReporter( - args.status, - repository, - start=start, - span=span, - label=label or repository, - **(stage_kwargs or {}), - ) - total_size = sum(size for _path, size, _checksum in selected) - downloaded = 0 - - def on_chunk(count): - nonlocal downloaded - downloaded += count - reporter.report(downloaded, total_size) - - for path, size, checksum in selected: - destination = Path(model_dir) / path - destination.parent.mkdir(parents=True, exist_ok=True) - _download_file_from_url( - url_builder(path), - destination, - size, - checksum, - progress=on_chunk, - retries=retries, - cancel_check=getattr(args, "cancel_check", None), - ) - - -def download_from_r2( - args, - model_dir, - start=0.05, - span=0.90, - label=None, - stage_kwargs=None, -): - """Download the declared model files from OmooLab R2.""" - _download_declared_files( - args, - model_dir, - lambda filename: r2_file_url(args.r2_directory, filename), - start, - span, - f"{label} (mirror)" if label else "mirror", - stage_kwargs, - ) - - -def download_from_hf( - args, - model_dir, - start=0.05, - span=0.90, - label=None, - stage_kwargs=None, -): - """Download the declared model files directly from Hugging Face resolve URLs.""" - _download_declared_files( - args, - model_dir, - lambda filename: hf_file_url(args.repository, filename), - start, - span, - label, - stage_kwargs, - retries=1, - ) - - -def format_rate(rate): - size = float(rate) - if size <= 0: - return "" - units = ("B/s", "KB/s", "MB/s", "GB/s") - unit_index = 0 - while size >= 1024 and unit_index < len(units) - 1: - size /= 1024 - unit_index += 1 - if unit_index == 0: - return f"{size:.0f} {units[unit_index]}" - return f"{size:.1f} {units[unit_index]}" - - -def infer_foreground( - input_path, - model_dir, - device, - model_cache=None, - cancel_check=None, -): - from PIL import Image - - input_path = Path(input_path) - check_cancelled(cancel_check) - if input_path.suffix.lower() not in IMAGE_SUFFIXES: - raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") - from ben2_onnx import create_session, infer - - owns_model = model_cache is None - requested_device = getattr(device, "type", device) - if owns_model: - load_started = time.perf_counter() - model = create_session(model_dir, requested_device) - load_ms = (time.perf_counter() - load_started) * 1000.0 - else: - model, load_ms = model_cache.get_ben2(model_dir, device) - - check_cancelled(cancel_check) - with Image.open(input_path) as opened_image: - image = opened_image.convert("RGB") - inference_started = time.perf_counter() - foreground = infer(model, image) - foreground = preserve_source_rgb(image, foreground) - inference_ms = (time.perf_counter() - inference_started) * 1000.0 - check_cancelled(cancel_check) - return foreground, load_ms, inference_ms - - -def preserve_source_rgb(source, foreground): - """Combine the untouched source RGB with BEN2's foreground alpha.""" - source_rgba = source.convert("RGBA") - alpha = foreground.convert("RGBA").getchannel("A") - if alpha.size != source_rgba.size: - from PIL import Image - - alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS) - source_rgba.putalpha(alpha) - return source_rgba - - -def infer_foreground_frames( - images, - model_dir, - device, - 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, - 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 Environment") - check_cancelled(cancel_check) - - output_directory = Path(args.output_dir) - output_directory.mkdir(parents=True, exist_ok=True) - device = args.device - - video_frames = getattr(args, "video_frames", 16) - start_frame = getattr(args, "start_frame", 1) - images, temporary_directory = prepare_background_input( - args.input, - video_frames, - output_directory, - start_frame=start_frame, - cancel_check=cancel_check, - ) - try: - write_status(args.status, 0.12, f"Loading BEN2 for {device.upper()}") - if len(images) == 1: - write_status(args.status, 0.35, "Removing background") - foreground, load_ms, inference_ms = infer_foreground( - images[0], - args.model_dir, - device, - 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, - 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{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, - 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": - scene_path = output_directory / SCENE_FILENAME - return [scene_path] if scene_path.is_file() else [] - if export_format == "gs_ply": - return sorted((output_directory / "gs_ply").glob("*.ply")) - if export_format == "depth_vis": - return sorted((output_directory / "depth_vis").glob("*.jpg")) - if export_format == "mesh": - result_path = output_directory / "mesh.glb" - return [result_path] if result_path.is_file() else [] - if export_format == "mesh_data": - result_path = output_directory / "mesh.npz" - return [result_path] if result_path.is_file() else [] - if export_format == "mini_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 [] - - -def npz_is_complete(path): - """Return True only when the npz archive can be fully read.""" - try: - import numpy as np - - with np.load(path, allow_pickle=False) as archive: - return "depth" in archive.files - except Exception: - return False - - -def wait_for_exported_paths( - output_directory, - export_format, - timeout=30.0, - cancel_check=None, -): - # DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整， - # 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取 - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - check_cancelled(cancel_check) - result_paths = exported_paths(output_directory, export_format) - if result_paths: - if ( - export_format in ("mini_npz", "npz") - and not npz_is_complete(result_paths[0]) - ): - time.sleep(0.1) - continue - return result_paths - time.sleep(0.05) - return [] - - -def move_exported_file( - source, - destination, - timeout=10.0, - cancel_check=None, -): - """Move an async export after Windows releases its file handle.""" - source = Path(source) - destination = Path(destination) - deadline = time.monotonic() + timeout - while True: - check_cancelled(cancel_check) - try: - source.replace(destination) - return destination - except OSError as error: - is_file_lock = isinstance(error, PermissionError) or getattr( - error, - "winerror", - None, - ) in {32, 33} - if not is_file_lock or time.monotonic() >= deadline: - raise - time.sleep(0.1) - - -def flatten_exported_paths( - output_directory, - result_paths, - cancel_check=None, -): - """Move files produced in exporter subdirectories to the job root.""" - output_directory = Path(output_directory) - flattened_paths = [] - for source in result_paths: - source = Path(source) - check_cancelled(cancel_check) - if source.parent == output_directory: - flattened_paths.append(source) - continue - destination = output_directory / source.name - if destination.exists(): - raise RuntimeError(f"Export filename already exists: {destination.name}") - flattened_paths.append( - move_exported_file( - source, - destination, - cancel_check=cancel_check, - ) - ) - for directory in sorted( - {path.parent for path in map(Path, result_paths)}, - key=lambda path: len(path.parts), - reverse=True, - ): - while directory != output_directory: - try: - directory.rmdir() - except OSError: - break - directory = directory.parent - return flattened_paths - - -def _save_moge_preview(output_directory, index, depth, normal, mask): - import numpy as np - from PIL import Image - - suffix = "" if index is None else f".{index + 1:04d}" - valid = mask & np.isfinite(depth) & (depth > 0) - preview = np.zeros(depth.shape, dtype=np.uint8) - if np.any(valid): - low, high = np.percentile(depth[valid], (2.0, 98.0)) - span = max(float(high - low), 1e-6) - preview[valid] = np.clip( - (depth[valid] - low) / span * 255.0, - 0.0, - 255.0, - ).astype(np.uint8) - Image.fromarray(preview, mode="L").save( - output_directory / f"depth{suffix}.png" - ) - normal_rgb = np.clip( - (normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0) - * 127.5, - 0.0, - 255.0, - ).astype(np.uint8) - Image.fromarray(normal_rgb, mode="RGB").save( - output_directory / f"normal{suffix}.png" - ) - Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save( - output_directory / f"mask{suffix}.png" - ) - - -def export_moge_debug_mesh( - output_directory, - points, - depth, - mask, - normal, - image, - export_format, -): - import numpy as np - import trimesh - - clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04) - faces = mesh_triangle_faces( - points, - depth, - np.ones(depth.shape, dtype=np.float32), - None, - False, - 0.0, - vertex_mask=clean_mask, - ) - vertices = points.reshape(-1, 3) - colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8) - vertex_normals = normal.reshape(-1, 3) - vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32) - vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32) - mesh = trimesh.Trimesh( - vertices=vertices, - faces=faces, - vertex_colors=colors, - vertex_normals=vertex_normals, - process=False, - ) - mesh.remove_unreferenced_vertices() - filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply" - mesh.export(Path(output_directory) / filename) - - -def export_moge_surface_archive( - output_path, - points, - depth, - mask, - normal, - intrinsics, - foreground_mask=None, -): - import numpy as np - - height, width = depth.shape - finite_mask = np.isfinite(depth) & (depth > 0.0) - finite_mask &= np.isfinite(points).all(axis=-1) - if foreground_mask is not None: - finite_mask &= foreground_mask - depth_discontinuity = depth_discontinuity_ratio(depth) - faces = mesh_triangle_faces( - points, - depth, - np.ones(depth.shape, dtype=np.float32), - None, - False, - 0.0, - vertex_mask=finite_mask, - ) - vertices = points.reshape(-1, 3) - u = np.linspace(0.0, 1.0, width, dtype=np.float32) - v = np.linspace(1.0, 0.0, height, dtype=np.float32) - uu, vv = np.meshgrid(u, v) - vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2) - pixel_intrinsics = intrinsics.copy() - pixel_intrinsics[0] *= width - pixel_intrinsics[1] *= height - pixel_intrinsics[2] = (0.0, 0.0, 1.0) - point_depth = depth.reshape(-1).astype(np.float32) - face_mask = mask.reshape(-1)[faces].all(axis=1) - face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) - masked_depth = point_depth[mask.reshape(-1).astype(bool)] - reference_depth = np.median(masked_depth if len(masked_depth) else point_depth) - np.savez( - output_path, - vertices=vertices.astype(np.float32), - faces=faces.astype(np.int32), - uv=vertex_uvs.astype(np.float32), - intrinsics=pixel_intrinsics.astype(np.float32), - image_size=np.array((width, height), dtype=np.int32), - reference_depth=np.float32(reference_depth), - point_depth=point_depth, - normal_map=normal.astype(np.float32), - face_mask=face_mask, - face_depth_discontinuity=face_depth_discontinuity, - ) - - -def run_moge(args, model_cache=None, cancel_check=None): - import numpy as np - from PIL import Image - from moge_onnx import ( - create_session, - infer as infer_moge_onnx, - provider_label, - ) - - write_status(args.status, 0.02, "Starting MoGe-2") - output_directory = Path(args.output_dir) - output_directory.mkdir(parents=True, exist_ok=True) - images, temporary_directory = prepare_input( - args.input, - args.video_frames, - output_directory, - start_frame=getattr(args, "start_frame", 1), - cancel_check=cancel_check, - ) - try: - if model_cache is None: - model = create_session(args.model_dir, args.device) - else: - model = model_cache.get_moge_onnx(args.model_dir, args.device) - write_status( - args.status, - 0.08, - f"Loading MoGe-2 on {provider_label(model)}", - ) - - depths = [] - normals = [] - masks = [] - points = [] - intrinsics = [] - source_images = [] - for index, image_path in enumerate(images): - check_cancelled(cancel_check) - write_status( - args.status, - 0.15 + 0.68 * index / len(images), - f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})", - ) - with Image.open(image_path) as opened: - image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 - source_images.append(image) - prediction = infer_moge_onnx( - model, - image, - int(args.resolution_level), - ) - depth = prediction["depth"] - normal = prediction["normal"] - mask = prediction["mask"].astype(bool) - point = prediction["points"] - intrinsic = prediction["intrinsics"] - depths.append(depth) - normals.append(normal) - masks.append(mask) - intrinsics.append(intrinsic) - if bool(getattr(args, "include_points", False)): - points.append(point) - debug_format = getattr(args, "debug_export_format", "") - if not debug_format or debug_format == "maps": - _save_moge_preview( - output_directory, - None if len(images) == 1 else index, - depth, - normal, - mask, - ) - - shapes = {array.shape for array in depths} - if len(shapes) != 1: - raise RuntimeError("MoGe-2 frame sizes must match") - result = { - "depth": np.stack(depths).astype(np.float32), - "normal": np.stack(normals).astype(np.float32), - "mask": np.stack(masks), - "intrinsics": np.stack(intrinsics).astype(np.float32), - } - if points: - result["points"] = np.stack(points).astype(np.float32) - debug_format = getattr(args, "debug_export_format", "") - if debug_format in {"glb", "ply"} and len(images) == 1: - write_status(args.status, 0.84, f"Exporting MoGe-2 {debug_format.upper()}") - export_moge_debug_mesh( - output_directory, - points[0], - depths[0], - masks[0], - normals[0], - source_images[0], - debug_format, - ) - - foreground_mask = None - remove_background_model_dir = getattr( - args, "remove_background_model_dir", "" - ) - if remove_background_model_dir: - write_status(args.status, 0.87, "Removing image background") - if len(images) == 1: - foreground, _, _ = infer_foreground( - images[0], - remove_background_model_dir, - args.device, - model_cache=model_cache, - cancel_check=cancel_check, - ) - foreground_path = output_directory / "foreground.png" - foreground.save(foreground_path) - if points: - foreground_mask = foreground_mesh_mask( - foreground_path, - result["depth"][0].shape, - ) - else: - infer_foreground_frames( - images, - remove_background_model_dir, - args.device, - output_directory, - args.status, - 0.87, - 0.08, - model_cache=model_cache, - cancel_check=cancel_check, - ) - if bool(getattr(args, "surface_mesh", False)): - if len(images) != 1 or not points: - raise RuntimeError("MoGe-2 Surface requires one point map") - write_status(args.status, 0.94, "Building official MoGe-2 surface") - export_moge_surface_archive( - output_directory / "mesh.npz", - points[0], - result["depth"][0], - result["mask"][0], - result["normal"][0], - result["intrinsics"][0], - foreground_mask=( - foreground_mask - if remove_background_model_dir - else None - ), - ) - np.savez(output_directory / "results.npz", **result) - copy_frames(images, output_directory) - write_status( - args.status, - 1.0, - f"Exported MoGe-2 data for {len(images)} frame(s)", - ) - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) - - -def merge_da3_frame_predictions(predictions): - """Combine independent single-frame DA3 results for existing exporters.""" - import numpy as np - - if not predictions: - raise ValueError("DA3 returned no frame predictions") - merged = predictions[0] - for name in ( - "depth", - "conf", - "extrinsics", - "intrinsics", - "processed_images", - ): - values = [getattr(prediction, name) for prediction in predictions] - setattr(merged, name, np.concatenate(values, axis=0)) - return merged - - -def reconstruct(args, model_cache=None, cancel_check=None): - write_status(args.status, 0.02, "Starting Environment") - check_cancelled(cancel_check) - from da3_onnx import create_session, infer - from onnx_runtime import provider_label - - output_directory = Path(args.output_dir) - output_directory.mkdir(parents=True, exist_ok=True) - write_status(args.status, 0.08, "Preparing input") - images, temporary_directory = prepare_input( - args.input, - args.video_frames, - output_directory, - start_frame=getattr(args, "start_frame", 1), - cancel_check=cancel_check, - ) - - try: - check_cancelled(cancel_check) - remove_background_model_dir = getattr( - args, - "remove_background_model_dir", - "", - ) - foreground_images = [] - if remove_background_model_dir: - write_status(args.status, 0.12, "Removing image background") - if len(images) == 1: - foreground, load_ms, inference_ms = infer_foreground( - images[0], - remove_background_model_dir, - args.device, - 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") - foreground_images = [output_directory / "foreground.png"] - else: - load_ms, inference_ms = infer_foreground_frames( - images, - remove_background_model_dir, - args.device, - output_directory, - args.status, - 0.12, - 0.05, - model_cache=model_cache, - cancel_check=cancel_check, - ) - foreground_images = sorted( - output_directory.glob("foreground.*.png") - ) - if remove_background_model_dir: - print( - f"BEN2 inference {inference_ms:.0f} ms; " - f"model load {load_ms:.0f} ms" - ) - write_status(args.status, 0.15, "Loading DA3 ONNX model") - if model_cache is None: - model = create_session(args.model_dir, args.device) - else: - model = model_cache.get_da3(args.model_dir, args.device) - write_status( - args.status, - 0.20, - f"Loaded DA3 on {provider_label(model)}", - ) - - check_cancelled(cancel_check) - write_status( - args.status, - 0.25, - f"Estimating depth for {len(images)} frame(s)", - ) - inference_options = { - "process_res": args.process_res, - "process_res_method": args.process_res_method, - "use_ray_pose": bool(args.use_ray_pose), - "ref_view_strategy": args.ref_view_strategy, - } - try: - if getattr(args, "inference_mode", "joint") == "per_frame": - predictions = [] - for index, image in enumerate(images): - check_cancelled(cancel_check) - write_status( - args.status, - 0.25 + 0.60 * index / len(images), - f"Estimating DA3 depth ({index + 1}/{len(images)})", - ) - predictions.append( - infer(model, [str(image)], **inference_options) - ) - prediction = merge_da3_frame_predictions(predictions) - else: - prediction = infer( - model, - [str(path) for path in images], - **inference_options, - ) - except UnicodeDecodeError as error: - # DirectML reports failures with a locale-specific message that - # ONNX Runtime cannot always decode, masking the real error. - if provider_label(model) == "DirectML": - raise RuntimeError( - "DA3 inference failed on DirectML. Switch the device to " - "CPU in SnapForm Preferences and retry. " - f"Original error: {error}" - ) from error - raise - check_cancelled(cancel_check) - write_status(args.status, 0.88, f"Exporting {args.export_format}") - if args.export_format == "mesh": - export_textured_mesh_glb( - prediction, - images, - output_directory, - getattr(args, "conf_thresh_percentile", 40.0), - bool(getattr(args, "mesh_filter_edges", True)), - float(getattr(args, "mesh_max_edge_length", 0.1)), - ) - elif args.export_format == "mesh_data": - def mesh_progress(step, total, label): - write_status( - args.status, - 0.88 + 0.09 * step / total, - f"Preparing surface data {step}/{total}: {label}", - ) - - export_image_mesh_data( - prediction, - images[0], - output_directory, - foreground_path=( - foreground_images[0] if foreground_images else None - ), - progress_callback=mesh_progress, - ) - elif args.export_format == "mini_npz": - export_da3_mini_npz(prediction, output_directory) - elif args.export_format == "depth_vis": - export_da3_depth_preview(prediction, output_directory) - elif args.export_format == "glb": - export_da3_point_cloud( - prediction, - output_directory, - args.conf_thresh_percentile, - args.num_max_points, - ) - else: - raise ValueError( - f"DA3 ONNX does not support export format: {args.export_format}" - ) - if args.export_format in {"glb", "gs_ply", "depth_vis", "mesh"}: - write_status( - args.status, - 0.97, - "Writing structured debug attributes", - ) - export_prediction_attributes(prediction, output_directory) - check_cancelled(cancel_check) - - result_paths = wait_for_exported_paths( - output_directory, - args.export_format, - cancel_check=cancel_check, - ) - if not result_paths: - raise RuntimeError( - f"No {args.export_format} files were exported" - ) - result_paths = flatten_exported_paths( - output_directory, - result_paths, - cancel_check=cancel_check, - ) - check_cancelled(cancel_check) - copy_frames(images, output_directory) - write_status( - args.status, - 1.0, - f"Exported {len(result_paths)} {args.export_format} file(s)", - ) - finally: - if temporary_directory is not None: - shutil.rmtree(temporary_directory, ignore_errors=True) - - -def parser(): - root = argparse.ArgumentParser() - commands = root.add_subparsers(dest="command", required=True) - - download = commands.add_parser("download") - download.add_argument("--repository", required=True) - download.add_argument("--model-dir", required=True) - download.add_argument("--status", required=True) - download.add_argument("--allow-pattern", action="append", default=[]) - download.add_argument("--required-pattern", action="append", default=[]) - download.add_argument("--r2-directory", default="") - download.add_argument("--r2-file", action="append", default=[]) - download.add_argument("--stage", type=int, default=None) - download.add_argument("--stages", type=int, default=None) - download.add_argument("--stage-label", default=None) - download.add_argument("--label", default=None) - download.add_argument("--progress-start", type=float, default=None) - download.add_argument("--progress-end", type=float, default=None) - - background = commands.add_parser("remove-background") - background.add_argument("--input", required=True) - background.add_argument("--model-dir", required=True) - background.add_argument("--output-dir", required=True) - background.add_argument("--status", required=True) - background.add_argument( - "--device", - choices=("auto", "cuda", "directml", "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) - reconstruction.add_argument("--model-dir", required=True) - reconstruction.add_argument("--output-dir", required=True) - reconstruction.add_argument( - "--export-format", - choices=EXPORT_FORMATS, - required=True, - ) - reconstruction.add_argument("--status", required=True) - reconstruction.add_argument("--remove-background-model-dir", default="") - reconstruction.add_argument("--process-res", type=int, required=True) - reconstruction.add_argument( - "--process-res-method", - choices=("upper_bound_resize", "lower_bound_resize"), - required=True, - ) - reconstruction.add_argument( - "--conf-thresh-percentile", - type=float, - required=True, - ) - reconstruction.add_argument("--num-max-points", type=int, required=True) - reconstruction.add_argument( - "--show-cameras", - type=int, - choices=(0, 1), - required=True, - ) - reconstruction.add_argument( - "--gs-views-interval", - type=int, - required=True, - ) - reconstruction.add_argument("--video-frames", type=int, required=True) - reconstruction.add_argument("--start-frame", type=int, default=1) - reconstruction.add_argument( - "--device", - choices=("auto", "cuda", "directml", "mps", "cpu"), - required=True, - ) - reconstruction.add_argument( - "--use-ray-pose", - type=int, - choices=(0, 1), - required=True, - ) - reconstruction.add_argument( - "--ref-view-strategy", - choices=("first", "middle", "saddle_balanced", "saddle_sim_range"), - required=True, - ) - return root - - -def main(): - args = parser().parse_args() - try: - if args.command == "download": - 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: - write_status(args.status, 1.0, "Task failed", str(error)) - raise - - -if __name__ == "__main__": - main() diff --git a/src/snapform/ai_service/server.py b/src/snapform/ai_service/server.py deleted file mode 100644 index c419f61..0000000 --- a/src/snapform/ai_service/server.py +++ /dev/null @@ -1,145 +0,0 @@ -import argparse -import sys -import threading -from pathlib import Path -from types import SimpleNamespace - - -SCRIPT_DIRECTORY = Path(__file__).resolve().parent -ADDON_DIRECTORY = SCRIPT_DIRECTORY.parent -for import_directory in (SCRIPT_DIRECTORY, ADDON_DIRECTORY): - if str(import_directory) not in sys.path: - sys.path.insert(0, str(import_directory)) - -import jobs as worker -from model_manager import ModelManager, ModelSessionCache -from blendjob.server import JobCancelled, JobServer, watch_parent - - -def create_server(models_directory=None, jobs_directory=None): - """Create the SnapForm-specific server and register its job types.""" - server = JobServer("SnapForm Job Server", jobs_directory=jobs_directory) - models = server.add_resource( - "models", - ModelManager(models_directory or Path.cwd() / "models"), - ) - - def run_worker(context, parameters, function, *, uses_models=True): - parameters = dict(parameters) - parameters["output_dir"] = str(context.directory) - models = context.resource("models") if uses_models else None - model_key = parameters.pop("model", "") - if model_key: - parameters["model_dir"] = str(models.directory(model_key)) - background_model = parameters.pop("remove_background_model", "") - if background_model: - parameters["remove_background_model_dir"] = str( - models.directory(background_model) - ) - args = SimpleNamespace(**parameters) - args.status = context.update - options = {"cancel_check": context.is_cancelled} - if uses_models: - options["model_cache"] = models - try: - function(args, **options) - except worker.TaskCancelled as error: - raise JobCancelled(str(error)) - - @server.job("remove-background") - def remove_background(context, parameters): - run_worker(context, parameters, worker.remove_background) - - @server.job("debug-ben2") - def debug_ben2(context, parameters): - run_worker(context, parameters, worker.remove_background) - - @server.job("extract-frames") - def extract_frames(context, parameters): - run_worker( - context, - parameters, - worker.extract_frames, - uses_models=False, - ) - - @server.job("generate-da3-depth") - def generate_da3_depth(context, parameters): - run_worker(context, parameters, worker.reconstruct) - - @server.job("generate-da3-surface") - def generate_da3_surface(context, parameters): - run_worker(context, parameters, worker.reconstruct) - - @server.job("debug-da3") - def debug_da3(context, parameters): - run_worker(context, parameters, worker.reconstruct) - - @server.job("generate-moge2-depth") - def generate_moge2_depth(context, parameters): - run_worker(context, parameters, worker.run_moge) - - @server.job("generate-moge2-surface") - def generate_moge2_surface(context, parameters): - parameters = {**parameters, "surface_mesh": True} - run_worker(context, parameters, worker.run_moge) - - @server.job("debug-moge2") - def debug_moge2(context, parameters): - parameters = dict(parameters) - parameters["debug_export_format"] = parameters.pop("export_format") - run_worker(context, parameters, worker.run_moge) - - @server.job("download-model") - def download_model(context, parameters): - context.resource("models").download(context, parameters["model"]) - - return server - - -def create_app(instance_id="test", server=None): - """Create the dedicated SnapForm FastAPI app.""" - server = server or create_server() - return server.create_app(instance_id) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--host", default="127.0.0.1") - parser.add_argument("--port", type=int, required=True) - parser.add_argument("--instance-id", required=True) - parser.add_argument("--parent-pid", type=int, required=True) - return parser.parse_args() - - -def main(): - import uvicorn - - args = parse_args() - job_server = create_server() - app = create_app(args.instance_id, job_server) - config = uvicorn.Config( - app, - host=args.host, - port=args.port, - access_log=False, - log_level="warning", - ) - http_server = uvicorn.Server(config) - - def stop_http_server(): - job_server.shutdown_event.wait() - http_server.should_exit = True - - threading.Thread(target=stop_http_server, daemon=True).start() - threading.Thread( - target=watch_parent, - args=(args.parent_pid, job_server.shutdown_event), - name="SnapFormParentMonitor", - daemon=True, - ).start() - http_server.run() - - -if __name__ == "__main__": - main() diff --git a/src/snapform/blender/depth_plane.py b/src/snapform/blender/depth_plane.py index 1611eb2..00a3b2e 100644 --- a/src/snapform/blender/depth_plane.py +++ b/src/snapform/blender/depth_plane.py @@ -1,5 +1,6 @@ import re import shutil +import tempfile from pathlib import Path import bpy @@ -53,7 +54,7 @@ VIDEO_NAME_SUFFIXES = ( def image_depth_model_items(owner, context): - from ..models import depth_plane_model_enum_items + from ..properties import depth_plane_model_enum_items return depth_plane_model_enum_items(owner, context) @@ -218,19 +219,17 @@ def process_resolution_scale_property(): def draw_process_resolution(layout, owner, image): process_res = scaled_process_resolution( image, - owner.snapform_process_res_factor, + owner.process_res_factor, ) layout.prop( owner, - "snapform_process_res_factor", + "process_res_factor", text=f"Resolution Scale ({process_res} px)", ) def save_packed_image(image, scene): - from ..client import temporary_input_path - - output_path = temporary_input_path() + output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png" image_settings = scene.render.image_settings original_file_format = image_settings.file_format try: @@ -286,9 +285,7 @@ def prepare_image_input(image, scene): return source_path, False sequence_frames = collect_sequence_frames(image) if sequence_frames: - from ..client import temporary_input_directory - - temporary_directory = Path(temporary_input_directory()) + temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-")) try: for index, frame_path in enumerate(sequence_frames): suffix = frame_path.suffix.lower() or ".jpg" @@ -1868,90 +1865,6 @@ def build_depth_plane( ) -class BuildDepthPlane(bpy.types.Operator): - bl_idname = "snapform.build_depth_plane" - bl_label = "Build Depth Plane" - bl_options = {"INTERNAL"} - - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_npz_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_depth_data_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - 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_delete_source: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_unshaded: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def execute(self, context): - source_object = bpy.data.objects.get(self.snapform_source_object_name) - if not is_image_empty(source_object): - self.report( - {"ERROR"}, - "The source image empty is no longer available", - ) - return {"CANCELLED"} - - npz_path = Path(self.snapform_npz_path) - if not npz_path.is_file(): - self.report( - {"ERROR"}, - f"Depth result does not exist: {npz_path}", - ) - return {"CANCELLED"} - - try: - foreground_image = None - if self.snapform_background_path: - foreground_image = create_removed_background_image( - source_object, - self.snapform_background_path, - assign_to_source=False, - ) - color_image = None - # The foreground result already carries the color, so the original - # frame sequence is unused when background removal is available. - if self.snapform_color_frames_dir and foreground_image is None: - color_image = create_sequence_image( - f"{image_base_name(source_object.data)}.frame", - self.snapform_color_frames_dir, - ) - build_depth_plane( - context, - source_object, - npz_path, - self.snapform_delete_source, - unshaded=self.snapform_unshaded, - depth_name=self.snapform_depth_data_name, - color_image=color_image, - foreground_image=foreground_image, - depth_frames_directory=( - self.snapform_depth_frames_dir or None - ), - ) - except (OSError, RuntimeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - def build_moge2_depth_plane( context, source_object, @@ -2048,177 +1961,6 @@ def build_moge2_depth_plane( ) -class BuildMoge2DepthPlane(bpy.types.Operator): - bl_idname = "snapform.build_moge2_depth_plane" - bl_label = "Build MoGe-2 Depth Plane" - bl_options = {"INTERNAL"} - - snapform_source_object_name: bpy.props.StringProperty(options={"HIDDEN"}) - snapform_npz_path: bpy.props.StringProperty(options={"HIDDEN"}) - snapform_background_path: bpy.props.StringProperty(options={"HIDDEN"}) - snapform_color_frames_dir: bpy.props.StringProperty(options={"HIDDEN"}) - snapform_data_frames_dir: bpy.props.StringProperty(options={"HIDDEN"}) - snapform_delete_source: bpy.props.BoolProperty(default=True, options={"HIDDEN"}) - snapform_unshaded: bpy.props.BoolProperty(default=False, options={"HIDDEN"}) - - def execute(self, context): - source_object = bpy.data.objects.get(self.snapform_source_object_name) - if not is_image_empty(source_object): - self.report({"ERROR"}, "The source image empty is no longer available") - return {"CANCELLED"} - try: - foreground_image = None - if self.snapform_background_path: - foreground_image = create_removed_background_image( - source_object, - self.snapform_background_path, - assign_to_source=False, - ) - color_image = None - # The foreground result already carries the color, so the original - # frame sequence is unused when background removal is available. - if self.snapform_color_frames_dir and foreground_image is None: - color_image = create_sequence_image( - f"{image_base_name(source_object.data)}.frame", - self.snapform_color_frames_dir, - ) - build_moge2_depth_plane( - context, - source_object, - Path(self.snapform_npz_path), - self.snapform_delete_source, - unshaded=self.snapform_unshaded, - color_image=color_image, - foreground_image=foreground_image, - data_frames_directory=( - self.snapform_data_frames_dir or None - ), - ) - except (OSError, RuntimeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - -class BuildOriginalDepthPlane(bpy.types.Operator): - bl_idname = "snapform.build_original_depth_plane" - bl_label = "Build Original Depth Plane" - bl_options = {"INTERNAL"} - - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - 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_depth_npz_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_delete_source: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_unshaded: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def execute(self, context): - source_object = bpy.data.objects.get(self.snapform_source_object_name) - if not is_image_empty(source_object): - self.report( - {"ERROR"}, - "The source image empty is no longer available", - ) - return {"CANCELLED"} - try: - foreground_image = None - if self.snapform_background_path: - foreground_image = create_removed_background_image( - source_object, - self.snapform_background_path, - assign_to_source=False, - ) - color_image = None - # The foreground result already carries the color, so the original - # frame sequence is unused when background removal is available. - if self.snapform_color_frames_dir and foreground_image is None: - color_image = create_sequence_image( - f"{image_base_name(source_object.data)}.frame", - self.snapform_color_frames_dir, - ) - depth_image = None - if self.snapform_depth_npz_path: - depth_arrays = extract_depth_arrays( - self.snapform_depth_npz_path - ) - if len(depth_arrays) == 1: - depth_image = create_depth_image( - depth_data_name(source_object), - depth_arrays[0], - ) - else: - depth_image = create_depth_sequence_image( - depth_data_name(source_object), - depth_arrays, - context.scene, - self.snapform_depth_frames_dir, - ) - build_image_depth_plane( - context, - source_object, - self.snapform_delete_source, - unshaded=self.snapform_unshaded, - color_image=color_image, - foreground_image=foreground_image, - depth_frames_directory=( - self.snapform_depth_frames_dir or None - ), - depth_image=depth_image, - ) - except (OSError, RuntimeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - -class ApplyRemovedBackground(bpy.types.Operator): - bl_idname = "snapform.apply_removed_background" - bl_label = "Apply Removed Background" - bl_options = {"INTERNAL"} - - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_output_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def execute(self, _context): - source_object = bpy.data.objects.get(self.snapform_source_object_name) - if not is_image_empty(source_object): - self.report( - {"ERROR"}, - "The source image empty is no longer available", - ) - return {"CANCELLED"} - try: - create_removed_background_image( - source_object, - self.snapform_output_path, - ) - except RuntimeError as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - class RemoveImageBackground(bpy.types.Operator): bl_idname = "snapform.remove_image_background" bl_label = "Remove Background" @@ -2227,15 +1969,13 @@ class RemoveImageBackground(bpy.types.Operator): @classmethod def poll(cls, context): - from ..operators import active_job - from ..client import service_busy, model_ready - from ..models import BEN2_MODEL + from ..job_runtime import runtime + from ..properties import BEN2_MODEL_KEY, model_ready return ( is_image_empty(context.object) - and active_job() is None - and not service_busy() - and model_ready(BEN2_MODEL.key) + and not runtime.server_busy() + and model_ready(BEN2_MODEL_KEY) ) def execute(self, context): @@ -2255,9 +1995,9 @@ class RemoveImageBackground(bpy.types.Operator): return {"CANCELLED"} result = bpy.ops.snapform.process_image_background( - snapform_input_path=str(source_path), - snapform_source_object_name=source_object.name, - snapform_delete_input=delete_input, + input_path=str(source_path), + source_object_name=source_object.name, + delete_input=delete_input, ) if "CANCELLED" not in result: return {"FINISHED"} @@ -2266,42 +2006,42 @@ class RemoveImageBackground(bpy.types.Operator): return {"CANCELLED"} -class ConvertImageToDepthPlane(bpy.types.Operator): - bl_idname = "snapform.convert_image_to_depth_plane" +class ConvertToDepthPlane(bpy.types.Operator): + bl_idname = "snapform.convert_to_depth_plane" bl_label = "Convert to Depth Plane" bl_description = "Convert an image to a displaced depth plane" bl_options = {"REGISTER"} - snapform_unshaded: bpy.props.BoolProperty( + unshaded: bpy.props.BoolProperty( name="Unshaded", description="Display the source color through an Emission shader", default=False, ) - snapform_remove_background: bpy.props.BoolProperty( + remove_background: bpy.props.BoolProperty( name="Remove Background", description="Remove the background during conversion with BEN2", default=False, ) - snapform_model: bpy.props.EnumProperty( + model: bpy.props.EnumProperty( name="Depth From", description="Model or original image used to build depth", items=image_depth_model_items, default=0, ) - snapform_use_ray_pose: bpy.props.BoolProperty( + use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", description="Legacy task compatibility; ONNX uses the camera decoder", default=False, ) - snapform_process_res_factor: process_resolution_scale_property() - snapform_moge_resolution_level: bpy.props.IntProperty( + process_res_factor: process_resolution_scale_property() + resolution_level: bpy.props.IntProperty( name="Resolution Level", description="MoGe-2 inference resolution level; higher is finer and slower", default=5, min=0, max=9, ) - snapform_keep_original: bpy.props.BoolProperty( + keep_original: bpy.props.BoolProperty( name="Keep Original", description="Keep the source Image Empty after conversion", default=False, @@ -2309,13 +2049,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator): @classmethod def poll(cls, context): - from ..operators import active_job - from ..client import service_busy + from ..job_runtime import runtime return ( is_image_empty(context.object) - and active_job() is None - and not service_busy() + and not runtime.server_busy() ) def invoke(self, context, _event): @@ -2323,15 +2061,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator): def draw(self, context): layout = self.layout - layout.prop(self, "snapform_model") - if self.snapform_model in {"SMALL", "BASE"}: + layout.prop(self, "model") + if self.model in {"SMALL", "BASE"}: draw_process_resolution(layout, self, context.object.data) - elif self.snapform_model.startswith("MOGE2_"): - layout.prop(self, "snapform_moge_resolution_level") - layout.prop(self, "snapform_remove_background") - layout.prop(self, "snapform_unshaded") + elif self.model.startswith("MOGE2_"): + layout.prop(self, "resolution_level") + layout.prop(self, "remove_background") + layout.prop(self, "unshaded") layout.separator(type="LINE") - layout.prop(self, "snapform_keep_original") + layout.prop(self, "keep_original") def execute(self, context): source_object = context.object @@ -2339,7 +2077,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): context.scene.snapform_settings.video_frames = ( default_extract_frames(source_object) ) - if self.snapform_model == "ORIGINAL": + if self.model == "ORIGINAL": try: source_path, delete_input = prepare_image_input( source_object.data, @@ -2350,15 +2088,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator): return {"CANCELLED"} run_original = ( bpy.ops.snapform.generate_original_background - if self.snapform_remove_background + if self.remove_background else bpy.ops.snapform.extract_original_frames ) result = run_original( - snapform_input_path=str(source_path), - snapform_source_object_name=source_object.name, - snapform_delete_source=not self.snapform_keep_original, - snapform_delete_input=delete_input, - snapform_unshaded=self.snapform_unshaded, + input_path=str(source_path), + source_object_name=source_object.name, + delete_source=not self.keep_original, + delete_input=delete_input, + unshaded=self.unshaded, ) if "CANCELLED" not in result: return {"FINISHED"} @@ -2366,7 +2104,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): remove_temporary_input(source_path) return {"CANCELLED"} - if self.snapform_model.startswith("MOGE2_"): + if self.model.startswith("MOGE2_"): try: source_path, delete_input = prepare_image_input( source_object.data, @@ -2376,16 +2114,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} result = bpy.ops.snapform.generate_moge2_depth( - snapform_input_path=str(source_path), - snapform_source_object_name=source_object.name, - snapform_moge_resolution_level=( - self.snapform_moge_resolution_level - ), - snapform_model_key=self.snapform_model, - snapform_delete_source=not self.snapform_keep_original, - snapform_delete_input=delete_input, - snapform_unshaded=self.snapform_unshaded, - snapform_remove_background=self.snapform_remove_background, + input_path=str(source_path), + source_object_name=source_object.name, + resolution_level=self.resolution_level, + model=self.model, + delete_source=not self.keep_original, + delete_input=delete_input, + unshaded=self.unshaded, + remove_background=self.remove_background, ) if "CANCELLED" not in result: return {"FINISHED"} @@ -2397,7 +2133,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): try: process_res = scaled_process_resolution( image, - self.snapform_process_res_factor, + self.process_res_factor, ) except ValueError as error: self.report({"ERROR"}, str(error)) @@ -2412,16 +2148,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator): return {"CANCELLED"} result = bpy.ops.snapform.generate_da3_depth( - snapform_input_path=str(source_path), - snapform_source_object_name=source_object.name, - snapform_depth_data_name=depth_data_name(source_object), - snapform_model_key=self.snapform_model, - snapform_use_ray_pose=self.snapform_use_ray_pose, - snapform_process_res=process_res, - snapform_delete_source=not self.snapform_keep_original, - snapform_delete_input=delete_input, - snapform_unshaded=self.snapform_unshaded, - snapform_remove_background=self.snapform_remove_background, + input_path=str(source_path), + source_object_name=source_object.name, + depth_data_name=depth_data_name(source_object), + model=self.model, + use_ray_pose=self.use_ray_pose, + process_res=process_res, + delete_source=not self.keep_original, + delete_input=delete_input, + unshaded=self.unshaded, + remove_background=self.remove_background, ) if "CANCELLED" in result: if delete_input: @@ -2435,11 +2171,11 @@ class SnapFormImageMenu(bpy.types.Menu): bl_label = "SnapForm" def draw(self, _context): - from ..client import env_installed + from ..job_runtime import runtime layout = self.layout layout.operator_context = "INVOKE_DEFAULT" - if not env_installed(): + if not runtime.environment_ready(): layout.operator( "snapform.install_environment", text="Install Environment", @@ -2447,12 +2183,12 @@ class SnapFormImageMenu(bpy.types.Menu): ) return layout.operator( - "snapform.convert_image_to_depth_plane", + "snapform.convert_to_depth_plane", text="Convert to Depth Plane", icon="MOD_DISPLACE", ) layout.operator( - "snapform.convert_image_to_mesh", + "snapform.convert_to_surface", text="Convert to Surface", icon="MESH_GRID", ) diff --git a/src/snapform/blender/surface.py b/src/snapform/blender/surface.py index e9f3844..0fba270 100644 --- a/src/snapform/blender/surface.py +++ b/src/snapform/blender/surface.py @@ -37,7 +37,7 @@ NODE_ASSET_NAME = "O_SnapForm.blend" def image_surface_model_items(owner, context): - from ..models import surface_model_enum_items + from ..properties import surface_model_enum_items return surface_model_enum_items(owner, context) @@ -554,96 +554,47 @@ def build_image_mesh( return mesh_object -class BuildImageMesh(bpy.types.Operator): - bl_idname = "snapform.build_image_mesh" - bl_label = "Build Image Surface" - bl_options = {"INTERNAL"} - - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_npz_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_delete_source: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_create_camera: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_unshaded: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_moge2: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def execute(self, context): - source_object = bpy.data.objects.get(self.snapform_source_object_name) - if not is_image_empty(source_object): - self.report({"ERROR"}, "The source Image Empty no longer exists") - return {"CANCELLED"} - try: - build_image_mesh( - context, - source_object, - self.snapform_npz_path, - self.snapform_delete_source, - self.snapform_create_camera, - self.snapform_unshaded, - self.snapform_moge2, - ) - except (OSError, RuntimeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - return {"FINISHED"} - - -class ConvertImageToMesh(bpy.types.Operator): - bl_idname = "snapform.convert_image_to_mesh" +class ConvertToSurface(bpy.types.Operator): + bl_idname = "snapform.convert_to_surface" bl_label = "Convert to Surface" bl_description = "Convert an Image Empty into a perspective textured surface" bl_options = {"REGISTER"} - snapform_model: bpy.props.EnumProperty( + model: bpy.props.EnumProperty( name="Depth From", description="Model used to estimate image depth", items=image_surface_model_items, default=0, ) - snapform_use_ray_pose: bpy.props.BoolProperty( + use_ray_pose: bpy.props.BoolProperty( name="Use Ray Pose", description="Legacy task compatibility; ONNX uses the camera decoder", default=False, ) - snapform_process_res_factor: process_resolution_scale_property() - snapform_moge_resolution_level: bpy.props.IntProperty( + process_res_factor: process_resolution_scale_property() + resolution_level: bpy.props.IntProperty( name="Resolution Level", description="MoGe-2 inference resolution level; higher is finer and slower", default=5, min=0, max=9, ) - snapform_create_camera: bpy.props.BoolProperty( + create_camera: bpy.props.BoolProperty( name="Create Matching Camera", description="Create a perspective camera matching the source projection", default=False, ) - snapform_unshaded: bpy.props.BoolProperty( + unshaded: bpy.props.BoolProperty( name="Unshaded", description="Display the source color through an Emission shader", default=False, ) - snapform_remove_background: bpy.props.BoolProperty( + remove_background: bpy.props.BoolProperty( name="Remove Background", description="Remove the background during conversion with BEN2", default=False, ) - snapform_keep_original: bpy.props.BoolProperty( + keep_original: bpy.props.BoolProperty( name="Keep Original", description="Keep the source Image Empty after conversion", default=False, @@ -651,18 +602,18 @@ class ConvertImageToMesh(bpy.types.Operator): @classmethod def poll(cls, context): - from ..models import MODELS, MOGE2_MODELS - from ..operators import active_job - from ..client import service_busy, model_ready + from ..job_runtime import runtime + from ..properties import model_catalog return ( is_image_empty(context.object) and video_total_frames(context.object.data) == 1 - and active_job() is None - and not service_busy() + and not runtime.server_busy() and any( - model_ready(model_key) - for model_key in (*MOGE2_MODELS, *MODELS) + model.get("ready") + and model.get("enabled") + and model.get("family") in {"moge2", "da3"} + for model in model_catalog() ) ) @@ -671,16 +622,16 @@ class ConvertImageToMesh(bpy.types.Operator): def draw(self, context): layout = self.layout - layout.prop(self, "snapform_model") - if self.snapform_model.startswith("MOGE2_"): - layout.prop(self, "snapform_moge_resolution_level") + layout.prop(self, "model") + if self.model.startswith("MOGE2_"): + layout.prop(self, "resolution_level") else: draw_process_resolution(layout, self, context.object.data) - layout.prop(self, "snapform_create_camera") - layout.prop(self, "snapform_remove_background") - layout.prop(self, "snapform_unshaded") + layout.prop(self, "create_camera") + layout.prop(self, "remove_background") + layout.prop(self, "unshaded") layout.separator(type="LINE") - layout.prop(self, "snapform_keep_original") + layout.prop(self, "keep_original") def execute(self, context): source_object = context.object @@ -689,10 +640,10 @@ class ConvertImageToMesh(bpy.types.Operator): return {"CANCELLED"} try: process_res = 1 - if not self.snapform_model.startswith("MOGE2_"): + if not self.model.startswith("MOGE2_"): process_res = scaled_process_resolution( source_object.data, - self.snapform_process_res_factor, + self.process_res_factor, ) source_path, delete_input = prepare_image_input( source_object.data, @@ -702,27 +653,25 @@ class ConvertImageToMesh(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} common = { - "snapform_input_path": str(source_path), - "snapform_source_object_name": source_object.name, - "snapform_model_key": self.snapform_model, - "snapform_delete_source": not self.snapform_keep_original, - "snapform_delete_input": delete_input, - "snapform_mesh_create_camera": self.snapform_create_camera, - "snapform_unshaded": self.snapform_unshaded, - "snapform_remove_background": self.snapform_remove_background, + "input_path": str(source_path), + "source_object_name": source_object.name, + "model": self.model, + "delete_source": not self.keep_original, + "delete_input": delete_input, + "create_camera": self.create_camera, + "unshaded": self.unshaded, + "remove_background": self.remove_background, } - if self.snapform_model.startswith("MOGE2_"): + if self.model.startswith("MOGE2_"): result = bpy.ops.snapform.generate_moge2_surface( **common, - snapform_moge_resolution_level=( - self.snapform_moge_resolution_level - ), + resolution_level=self.resolution_level, ) else: result = bpy.ops.snapform.generate_da3_surface( **common, - snapform_use_ray_pose=self.snapform_use_ray_pose, - snapform_process_res=process_res, + use_ray_pose=self.use_ray_pose, + process_res=process_res, ) if "CANCELLED" in result: if delete_input: diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index b55382e..79f5127 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -12,7 +12,7 @@ platforms = ["windows-x64", "macos-arm64", "linux-x64"] license = ["SPDX:GPL-3.0-or-later"] [permissions] -network = "Access the local Service and download environment dependencies and model weights" +network = "Access the local Job Server and download environment dependencies and model weights" files = "Read source images and store environment data" [build] diff --git a/src/snapform/blendjob/__init__.py b/src/snapform/blendjob/__init__.py index 914b8d9..4f21143 100644 --- a/src/snapform/blendjob/__init__.py +++ b/src/snapform/blendjob/__init__.py @@ -1,37 +1,21 @@ -"""Reusable local HTTP job service for Blender add-ons.""" +"""Reusable local HTTP Job Server for Blender add-ons.""" -from .client import JobClient, JobRequest +from .client import JobClient, JobResult +from .controller import ServerConnection, ServerController +from .operator import JobOperatorBase, JobOperatorState +from .runtime import BlenderJobRuntime from .server import JobCancelled, JobContext, JobServer -from .service import ServiceConnection, ServiceController - -try: - from .operator import ( - ModalJobOperatorBase, - ModalJobState, - active_job, - cancel_active_job, - close_active_job, - ) -except ModuleNotFoundError as error: - if error.name != "bpy": - raise __all__ = ( "JobCancelled", "JobClient", "JobContext", - "JobRequest", + "JobResult", "JobServer", - "ServiceConnection", - "ServiceController", + "BlenderJobRuntime", + "JobOperatorBase", + "JobOperatorState", + "ServerConnection", + "ServerController", ) - -if "ModalJobOperatorBase" in globals(): - __all__ += ( - "ModalJobOperatorBase", - "ModalJobState", - "active_job", - "cancel_active_job", - "close_active_job", - ) diff --git a/src/snapform/blendjob/client.py b/src/snapform/blendjob/client.py index d407ff1..a65d5af 100644 --- a/src/snapform/blendjob/client.py +++ b/src/snapform/blendjob/client.py @@ -1,5 +1,6 @@ import json from dataclasses import dataclass +from pathlib import Path from urllib.error import HTTPError, URLError from urllib.request import ProxyHandler, Request, build_opener @@ -42,19 +43,13 @@ class JobClient: def health(self): return self.request("GET", "/health") - def submit(self, job_type, **parameters): + def submit(self, job_type, parameters): return self.request( "POST", "/jobs", - { - "job_type": job_type, - "parameters": parameters, - }, + {"job_type": str(job_type), "parameters": dict(parameters)}, ) - def submit_request(self, request): - return self.request("POST", "/jobs", request.payload) - def status(self, job_id): return self.request("GET", f"/jobs/{job_id}") @@ -72,28 +67,35 @@ class JobClient: def shutdown(self): return self.request("POST", "/shutdown") -@dataclass(frozen=True, init=False) -class JobRequest: - """Serializable request created before a Job is submitted.""" +@dataclass(frozen=True) +class JobResult: + """Successful Job value and files returned to a Blender Operator.""" - job_type: str - parameters: dict - - def __init__(self, job_type, parameters=None): - if isinstance(job_type, dict): - payload = job_type - job_type = payload.get("job_type", "") - parameters = payload.get("parameters", {}) - object.__setattr__(self, "job_type", str(job_type)) - object.__setattr__(self, "parameters", dict(parameters or {})) + job_id: str + directory: Path + value: object + status: dict @classmethod - def create(cls, job_type, **parameters): - return cls(job_type, parameters) + def from_status(cls, status): + return cls( + job_id=str(status.get("job_id", "")), + directory=Path(status["directory"]), + value=status.get("result"), + status=dict(status), + ) - @property - def payload(self): - return { - "job_type": self.job_type, - "parameters": self.parameters, - } + def file(self, name): + if not isinstance(self.value, dict): + raise RuntimeError("Job result does not contain named files") + try: + relative = self.value[name] + except KeyError: + raise RuntimeError(f"Job result does not contain file: {name}") from None + root = self.directory.resolve() + path = (root / str(relative)).resolve() + if path != root and root not in path.parents: + raise RuntimeError(f"Job result file escapes its directory: {relative}") + if not path.is_file(): + raise RuntimeError(f"Job result file does not exist: {path}") + return path diff --git a/src/snapform/blendjob/service.py b/src/snapform/blendjob/controller.py similarity index 75% rename from src/snapform/blendjob/service.py rename to src/snapform/blendjob/controller.py index dd22ddf..d509b53 100644 --- a/src/snapform/blendjob/service.py +++ b/src/snapform/blendjob/controller.py @@ -4,18 +4,18 @@ import time import uuid from dataclasses import dataclass -from .client import JobClient +from .client import JobClient, JobResult @dataclass -class ServiceConnection: +class ServerConnection: port: int pid: int instance_id: str client: JobClient -class ServiceController: +class ServerController: """Own a dedicated local JobServer process and its HTTP connection.""" def __init__( @@ -61,7 +61,7 @@ class ServiceController: self.log_file = None raise client = JobClient(self.host, port) - self.connection = ServiceConnection( + self.connection = ServerConnection( port, self.process.pid, instance_id, @@ -72,10 +72,10 @@ class ServiceController: def health(self, timeout=0.3): if self.connection is None: - raise RuntimeError("Service is not running") + raise RuntimeError("Server is not running") health = self.connection.client.request("GET", "/health", timeout=timeout) if health.get("instance_id") != self.connection.instance_id: - raise RuntimeError("Service instance does not match") + raise RuntimeError("Server instance does not match") state = "BUSY" if health.get("busy") else "READY" self.snapshot = { **health, @@ -97,7 +97,7 @@ class ServiceController: if self.process is None: self.start() time.sleep(0.05) - raise RuntimeError(f"Service did not become ready; log: {self.log_path_factory()}") + raise RuntimeError(f"Server did not become ready; log: {self.log_path_factory()}") def enable(self): self.auto_start = True @@ -129,8 +129,51 @@ class ServiceController: self.auto_start = True return self.start() - def submit(self, request): - return self.ensure().client.submit_request(request) + def submit(self, job_type, parameters): + return self.ensure().client.submit(job_type, parameters) + + def request( + self, + job_type, + parameters, + *, + response=None, + progress=None, + cancel_check=None, + poll_interval=0.25, + ): + """Submit one Job and return its response after polling to completion.""" + submitted = self.submit(job_type, parameters) + job_id = submitted["job_id"] + directory = submitted["directory"] + cancel_requested = False + try: + while True: + if ( + not cancel_requested + and cancel_check is not None + and cancel_check() + ): + self.cancel(job_id) + cancel_requested = True + status = self.status(job_id) + if progress is not None: + progress({"job_id": job_id, "directory": directory, **status}) + state = status.get("state") + if state == "succeeded": + result = JobResult.from_status( + {"job_id": job_id, "directory": directory, **status} + ) + if response is not None: + response(result) + return result + if state == "cancelled": + raise RuntimeError("Server request was cancelled") + if state == "failed": + raise RuntimeError(status.get("error") or "Server request failed") + time.sleep(max(float(poll_interval), 0.01)) + finally: + self.mark_job_complete(job_id) def status(self, job_id): return self.ensure().client.status(job_id) diff --git a/src/snapform/blendjob/environment.py b/src/snapform/blendjob/environment.py new file mode 100644 index 0000000..a14346c --- /dev/null +++ b/src/snapform/blendjob/environment.py @@ -0,0 +1,351 @@ +import argparse +import hashlib +import json +import os +import platform +import queue +import shutil +import subprocess +import sys +import tarfile +import tempfile +import threading +import time +import urllib.request +import zipfile +from pathlib import Path + + +UV_VERSION = "0.11.3" +ENVIRONMENT_SCHEMA = 1 +BASE_PACKAGES = ( + "fastapi==0.139.2", + "uvicorn==0.51.0", +) +UV_ARTIFACTS = { + "x86_64-pc-windows-msvc": ( + "uv-x86_64-pc-windows-msvc.zip", + "ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c", + ), + "aarch64-apple-darwin": ( + "uv-aarch64-apple-darwin.tar.gz", + "2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f", + ), + "x86_64-unknown-linux-gnu": ( + "uv-x86_64-unknown-linux-gnu.tar.gz", + "c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce", + ), +} +UV_SOURCES = ( + f"https://github.com/astral-sh/uv/releases/download/{UV_VERSION}", + f"https://mirror.omoolab.xyz/uv/{UV_VERSION}", +) + + +def normalized_environment(environment): + """Validate and normalize a public Environment dictionary.""" + if not isinstance(environment, dict): + raise TypeError("environment must be a dict") + python = str(environment.get("python", "")).strip() + if not python: + raise ValueError("environment['python'] is required") + packages = environment.get("packages", ()) + platform_packages = environment.get("platform_packages", {}) + if not isinstance(packages, (list, tuple)): + raise TypeError("environment['packages'] must be a list") + if not isinstance(platform_packages, dict): + raise TypeError("environment['platform_packages'] must be a dict") + normalized_platforms = {} + for name, values in platform_packages.items(): + if not isinstance(values, (list, tuple)): + raise TypeError(f"platform_packages[{name!r}] must be a list") + normalized_platforms[str(name)] = [str(value) for value in values] + return { + "python": python, + "packages": [str(value) for value in packages], + "platform_packages": normalized_platforms, + } + + +def environment_digest(environment): + payload = json.dumps( + { + "schema": ENVIRONMENT_SCHEMA, + "base_packages": BASE_PACKAGES, + "environment": normalized_environment(environment), + }, + sort_keys=True, + separators=(",", ":"), + ).encode("utf-8") + return hashlib.sha256(payload).hexdigest() + + +def platform_name(): + if sys.platform == "win32": + return "windows" + if sys.platform == "darwin": + return "macos" + return "linux" + + +def resolved_packages(environment): + environment = normalized_environment(environment) + platform_packages = environment["platform_packages"] + selected = platform_packages.get( + platform_name(), + platform_packages.get("default", ()), + ) + return [*BASE_PACKAGES, *environment["packages"], *selected] + + +def environment_python(storage_root): + directory = Path(storage_root) / ".venv" + if os.name == "nt": + return directory / "Scripts" / "python.exe" + return directory / "bin" / "python" + + +def uv_path(storage_root): + filename = "uv.exe" if os.name == "nt" else "uv" + return Path(storage_root) / "tools" / filename + + +def uv_target(): + machine = platform.machine().lower() + if sys.platform == "win32" and machine in {"amd64", "x86_64"}: + return "x86_64-pc-windows-msvc" + if sys.platform == "darwin" and machine in {"arm64", "aarch64"}: + return "aarch64-apple-darwin" + if sys.platform.startswith("linux") and machine in {"amd64", "x86_64"}: + return "x86_64-unknown-linux-gnu" + raise RuntimeError(f"Unsupported uv platform: {sys.platform} {machine}") + + +def write_status(path, progress, message, *, stage, stages, error=None): + payload = { + "progress": min(max(float(progress), 0.0), 1.0), + "message": str(message), + "stage": int(stage), + "stages": int(stages), + "stage_label": "Environment", + } + if error: + payload["error"] = str(error) + Path(path).write_text( + json.dumps(payload, ensure_ascii=False), + encoding="utf-8", + ) + + +def child_environment(): + environment = os.environ.copy() + environment["PYTHONIOENCODING"] = "utf-8:backslashreplace" + environment["PYTHONUTF8"] = "1" + for name, value in urllib.request.getproxies().items(): + environment.setdefault(f"{name.upper()}_PROXY", value) + return environment + + +def _download(url, destination, status, stage, stages, start, end): + request = urllib.request.Request(url, headers={"User-Agent": "BlendJob"}) + with urllib.request.urlopen(request) as response, destination.open("wb") as output: + total = int(response.headers.get("Content-Length") or 0) + received = 0 + while True: + chunk = response.read(256 * 1024) + if not chunk: + break + output.write(chunk) + received += len(chunk) + fraction = received / total if total else 0.5 + progress = start + (end - start) * min(fraction, 1.0) + write_status( + status, + progress, + "Installing uv", + stage=stage, + stages=stages, + ) + + +def _verify(path, checksum): + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + if digest.hexdigest().lower() != checksum.lower(): + raise RuntimeError("uv archive checksum does not match") + + +def _extract_uv(archive, destination): + destination.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory() as directory: + extracted = Path(directory) + if archive.name.endswith(".zip"): + with zipfile.ZipFile(archive) as package: + package.extractall(extracted) + else: + with tarfile.open(archive, "r:gz") as package: + package.extractall(extracted) + candidates = tuple(extracted.rglob(destination.name)) + if not candidates: + raise RuntimeError("uv executable was not found in its archive") + shutil.copy2(candidates[0], destination) + if os.name != "nt": + destination.chmod(destination.stat().st_mode | 0o111) + + +def install_uv(storage_root, status, stage, stages, start=0.0, end=0.15): + destination = uv_path(storage_root) + if destination.is_file(): + write_status(status, end, "uv is ready", stage=stage, stages=stages) + return destination + target = uv_target() + filename, checksum = UV_ARTIFACTS[target] + errors = [] + destination.parent.mkdir(parents=True, exist_ok=True) + for source in UV_SOURCES: + descriptor, archive_name = tempfile.mkstemp(suffix=f"-{filename}") + os.close(descriptor) + archive = Path(archive_name) + try: + _download( + f"{source}/{filename}", + archive, + status, + stage, + stages, + start, + end * 0.8, + ) + _verify(archive, checksum) + _extract_uv(archive, destination) + write_status(status, end, "uv is ready", stage=stage, stages=stages) + return destination + except Exception as error: + errors.append(f"{source}: {error}") + finally: + archive.unlink(missing_ok=True) + raise RuntimeError("Unable to install uv: " + "; ".join(errors)) + + +def run_command(command, environment=None): + subprocess.check_call( + [str(value) for value in command], + env=environment or child_environment(), + ) + + +def install_packages( + command, + status, + stage, + stages, + start=0.22, + end=0.98, + clock=time.monotonic, +): + process = subprocess.Popen( + [str(value) for value in command], + env=child_environment(), + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + encoding="utf-8", + errors="replace", + ) + output = queue.Queue() + + def read_output(): + for line in process.stdout: + output.put(line) + output.put(None) + + threading.Thread(target=read_output, daemon=True).start() + started = clock() + last_line = "Installing packages" + while True: + try: + line = output.get(timeout=0.2) + except queue.Empty: + line = "" + if line is None: + break + if line.strip(): + last_line = line.strip() + print(last_line) + elapsed = clock() - started + fraction = min(elapsed / (elapsed + 30.0), 0.97) + write_status( + status, + start + (end - start) * fraction, + last_line, + stage=stage, + stages=stages, + ) + return_code = process.wait() + if return_code: + raise subprocess.CalledProcessError(return_code, command) + + +def install(storage_root, environment, status, stages=1): + environment = normalized_environment(environment) + storage_root = Path(storage_root).resolve() + storage_root.mkdir(parents=True, exist_ok=True) + venv = storage_root / ".venv" + manifest = storage_root / "manifest.json" + expected_digest = environment_digest(environment) + uv = install_uv(storage_root, status, 1, stages) + manifest.unlink(missing_ok=True) + write_status(status, 0.16, "Creating Python Environment", stage=1, stages=stages) + run_command([uv, "venv", "--clear", "--python", environment["python"], venv]) + write_status(status, 0.22, "Installing packages", stage=1, stages=stages) + packages = resolved_packages(environment) + install_packages( + [uv, "pip", "install", "--python", environment_python(storage_root), *packages], + status, + 1, + stages, + ) + manifest.write_text( + json.dumps( + { + "environment_hash": expected_digest, + "python": environment["python"], + "packages": packages, + }, + indent=2, + ), + encoding="utf-8", + ) + write_status(status, 1.0, "Environment is ready", stage=1, stages=stages) + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--storage-root", required=True) + parser.add_argument("--config", required=True) + parser.add_argument("--status", required=True) + parser.add_argument("--stages", type=int, default=1) + return parser.parse_args() + + +def main(): + args = parse_args() + try: + environment = json.loads(Path(args.config).read_text(encoding="utf-8")) + install(args.storage_root, environment, args.status, args.stages) + except Exception as error: + write_status( + args.status, + 1.0, + "Environment installation failed", + stage=1, + stages=args.stages, + error=error, + ) + raise + + +if __name__ == "__main__": + main() diff --git a/src/snapform/blendjob/installer/environment.py b/src/snapform/blendjob/installer/environment.py new file mode 100644 index 0000000..e5716ad --- /dev/null +++ b/src/snapform/blendjob/installer/environment.py @@ -0,0 +1,13 @@ +import sys +from pathlib import Path + + +BLENDJOB_DIRECTORY = Path(__file__).resolve().parent.parent +if str(BLENDJOB_DIRECTORY) not in sys.path: + sys.path.insert(0, str(BLENDJOB_DIRECTORY)) + +from environment import main + + +if __name__ == "__main__": + main() diff --git a/src/snapform/blendjob/operator.py b/src/snapform/blendjob/operator.py index 65bfb24..2366ce6 100644 --- a/src/snapform/blendjob/operator.py +++ b/src/snapform/blendjob/operator.py @@ -1,19 +1,18 @@ +import json import threading from dataclasses import dataclass from pathlib import Path -import bpy - -from .client import JobRequest +from .client import JobResult @dataclass -class ModalJobState: - """Runtime state shared by all Blender modal Job operators.""" +class JobOperatorState: + """Internal state for one modal Job Operator invocation.""" + runtime: object + operator: object controller: object - data: object = None - cleanup: object = None job_id: str = "" directory: Path | None = None started: bool = False @@ -23,121 +22,110 @@ class ModalJobState: progress_stage: int | None = None progress: float = 0.0 - def __getattr__(self, name): - if name == "result_directory": - return self.directory - if self.data is None: - raise AttributeError(name) - return getattr(self.data, name) + +def operator_properties(operator): + """Return every custom Blender Property declared by an Operator class.""" + names = [] + for class_type in reversed(type(operator).__mro__): + for name in getattr(class_type, "__annotations__", {}): + if name not in names: + names.append(name) + parameters = { + name: _json_value(getattr(operator, name)) + for name in names + if hasattr(operator, name) + } + json.dumps(parameters) + return parameters -_active_job = None - - -def active_job(): - return _active_job - - -def cancel_active_job(): - job = active_job() - if job is None: - return - job.cancelled = True - if not job.started or not job.job_id: - return +def _json_value(value): + if value is None or isinstance(value, (bool, float, int, str)): + return value + if isinstance(value, (list, tuple)): + return [_json_value(item) for item in value] try: - job.controller.cancel(job.job_id) - except RuntimeError: - pass + return [_json_value(item) for item in value] + except TypeError: + return value -def close_active_job(): - global _active_job - - if _active_job is not None: - if _active_job.job_id: - _active_job.controller.mark_job_complete(_active_job.job_id) - if _active_job.cleanup is not None: - _active_job.cleanup(_active_job.data) - _active_job.cleanup = None - _active_job = None - - -class ModalJobOperatorBase(bpy.types.Operator): - """Ready-to-use Blender lifecycle for one HTTP Job request.""" +class JobOperatorBase: + """Blender mixin that turns one Operator into a complete remote Job.""" bl_options = {"INTERNAL"} poll_interval = 0.25 starting_message = "Starting task" - controller_factory = None + job_runtime = None job_type = "" - def get_service_controller(self): - """Return the ServiceController used by this add-on.""" - if self.controller_factory is None: - raise RuntimeError("Set controller_factory on the Job operator") - return self.controller_factory() + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + annotations = {} + for base in reversed(cls.__mro__[1:]): + annotations.update(getattr(base, "__annotations__", {})) + annotations.update(cls.__dict__.get("__annotations__", {})) + if annotations: + cls.__annotations__ = annotations - def create_job_request(self, context): - """Create the JobRequest for this operator invocation.""" - if not self.job_type: - raise RuntimeError("Set job_type on the Job operator") - return JobRequest(self.job_type, self.job_parameters(context)) + def request(self, _context): + """Return the complete JSON parameters submitted to the Server.""" + return operator_properties(self) - def job_parameters(self, _context): - """Return JSON-compatible business parameters for this invocation.""" - return {} + def response(self, _context, _result): + """Apply an optional successful result on Blender's main thread.""" - def create_job_data(self, _context): - """Return optional plug-in data retained until completion.""" - return None + def cleanup(self): + """Release invocation resources after every terminal outcome.""" + + def controller(self, runtime): + """Return the operation controller used by this invocation.""" + return runtime.server def execute(self, context): - global _active_job - - if active_job() is not None: + runtime = self._runtime() + if runtime.active_job is not None: self.report({"WARNING"}, "Another task is running") return {"CANCELLED"} + if not self.job_type: + self.report({"ERROR"}, "Set job_type on the Job Operator") + return {"CANCELLED"} try: - request = self.create_job_request(context) - if not isinstance(request, JobRequest): - raise TypeError("create_job_request() must return JobRequest") - controller = self.get_service_controller() - data = self.create_job_data(context) + parameters = self.request(context) + if not isinstance(parameters, dict): + raise TypeError("request(context) must return a dict") + json.dumps(parameters) + controller = self.controller(runtime) except (OSError, RuntimeError, TypeError, ValueError) as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} - job = ModalJobState( - controller=controller, - data=data, - cleanup=self.cleanup_job_data, - ) - _active_job = job + job = JobOperatorState(runtime, self, controller) + runtime.begin_job(job) if hasattr(context.window_manager, "progress_begin"): context.window_manager.progress_begin(0, 100) - self.update_job_ui(context, 0.0, self.get_starting_message(context)) + runtime.update_ui(context, 0.0, self.starting_message) self._timer = context.window_manager.event_timer_add( self.poll_interval, window=context.window, ) context.window_manager.modal_handler_add(self) - self.redraw_job_ui(context, force=True) + runtime.redraw_ui(context, force=True) job.start_thread = threading.Thread( target=self._submit_job, - args=(job, request), + args=(job, parameters), name="BlendJobSubmit", daemon=True, ) job.start_thread.start() return {"RUNNING_MODAL"} - def _submit_job(self, job, request): + def _submit_job(self, job, parameters): if job.cancelled: job.started = True return try: - submitted = job.controller.submit(request) + submitted = job.controller.submit(self.job_type, parameters) job.job_id = submitted["job_id"] job.directory = Path(submitted["directory"]) job.started = True @@ -147,53 +135,54 @@ class ModalJobOperatorBase(bpy.types.Operator): job.start_error = str(error) def modal(self, context, event): + runtime = self._runtime() if event.type == "ESC": - cancel_active_job() + runtime.cancel_active() if event.type != "TIMER": return {"PASS_THROUGH"} - job = active_job() + job = runtime.active_job if job is None: self.remove_job_timer(context) return {"CANCELLED"} if job.start_error: - return self._finish_start_error(context, job) + self.report({"ERROR"}, f"Unable to start task: {job.start_error}") + return self._close(context, job, "Task failed to start", cancelled=True) if not job.started: return {"RUNNING_MODAL"} + if job.cancelled and not job.job_id: + return self._close(context, job, "Task cancelled", cancelled=True) try: status = job.controller.status(job.job_id) except RuntimeError as error: - return self._finish_failure(context, job, str(error)) + return self._finish_failure(context, job, error) self.update_from_job_status(context, job, status) state = status.get("state") if state == "succeeded": return self._finish_success(context, job, status) if state in {"failed", "cancelled"}: - return self._finish_terminal(context, job, status) + if job.cancelled or state == "cancelled": + return self._close(context, job, "Task cancelled", cancelled=True) + return self._finish_failure( + context, + job, + status.get("error") or "Server task failed", + ) return {"RUNNING_MODAL"} - def _finish_start_error(self, context, job): - self.report({"ERROR"}, f"Unable to start task: {job.start_error}") - return self._close(context, job, "Task failed to start", cancelled=True) - - def _finish_terminal(self, context, job, status): - if job.cancelled or status.get("state") == "cancelled": - return self._close(context, job, "Task cancelled", cancelled=True) - return self._finish_failure( - context, - job, - status.get("error") or "Service task failed", - ) - - def _finish_failure(self, context, job, error): - self.report({"ERROR"}, str(error)) - return self._close(context, job, "Task failed", cancelled=True) - def _finish_success(self, context, job, status): try: - message = self.on_job_succeeded(context, job, status) - except (OSError, RuntimeError, ValueError) as error: + complete_status = { + "job_id": job.job_id, + "directory": str(job.directory), + **status, + } + message = self.response( + context, + JobResult.from_status(complete_status), + ) + except (OSError, RuntimeError, TypeError, ValueError) as error: return self._finish_failure(context, job, error) return self._close( context, @@ -202,18 +191,19 @@ class ModalJobOperatorBase(bpy.types.Operator): cancelled=False, ) + def _finish_failure(self, context, job, error): + self.report({"ERROR"}, str(error)) + return self._close(context, job, "Task failed", cancelled=True) + def _close(self, context, job, message, *, cancelled): self.remove_job_timer(context) if hasattr(context.window_manager, "progress_end"): context.window_manager.progress_end() - close_active_job() - self.update_job_ui(context, 0.0 if cancelled else 1.0, message) - self.redraw_job_ui(context, force=True) + job.runtime.finish_job(job) + job.runtime.update_ui(context, 0.0 if cancelled else 1.0, message) + job.runtime.redraw_ui(context, force=True) return {"CANCELLED"} if cancelled else {"FINISHED"} - def get_starting_message(self, _context): - return self.starting_message - def update_from_job_status(self, context, job, status): reported = min(max(float(status.get("progress", job.progress)), 0.0), 1.0) stage = status.get("stage") @@ -223,30 +213,8 @@ class ModalJobOperatorBase(bpy.types.Operator): else: job.progress = max(job.progress, reported) message = self.format_job_status(status, self.starting_message) - self.update_job_ui(context, job.progress, message) - self.redraw_job_ui(context, force=True) - - def update_job_ui(self, context, progress, _message): - """Update Blender's standard progress indicator by default.""" - if hasattr(context.window_manager, "progress_update"): - context.window_manager.progress_update(int(progress * 100)) - - def on_job_succeeded(self, _context, _job, status): - """Import results if needed and optionally return a final message.""" - return status.get("message") - - def cleanup_job_data(self, _data): - pass - - def redraw_job_ui(self, _context, force=False): - del force - - def remove_job_timer(self, context): - timer = getattr(self, "_timer", None) - if timer is None: - return - context.window_manager.event_timer_remove(timer) - self._timer = None + job.runtime.update_ui(context, job.progress, message) + job.runtime.redraw_ui(context, force=True) def format_job_status(self, status, fallback): message = status.get("message") @@ -261,7 +229,27 @@ class ModalJobOperatorBase(bpy.types.Operator): return f"{prefix} {stage_label}" if stage_label else prefix return str(message or stage_label or fallback).strip() + def remove_job_timer(self, context): + timer = getattr(self, "_timer", None) + if timer is None: + return + context.window_manager.event_timer_remove(timer) + self._timer = None + def cancel(self, context): - cancel_active_job() + runtime = self._runtime() + job = runtime.active_job + runtime.cancel_active() self.remove_job_timer(context) - self.redraw_job_ui(context, force=True) + if job is None: + return + if hasattr(context.window_manager, "progress_end"): + context.window_manager.progress_end() + runtime.finish_job(job) + runtime.update_ui(context, 0.0, "Task cancelled") + runtime.redraw_ui(context, force=True) + + def _runtime(self): + if self.job_runtime is None: + raise RuntimeError("Use BlenderJobRuntime.JobOperatorBase") + return self.job_runtime diff --git a/src/snapform/blendjob/runner.py b/src/snapform/blendjob/runner.py new file mode 100644 index 0000000..480b87f --- /dev/null +++ b/src/snapform/blendjob/runner.py @@ -0,0 +1,71 @@ +import argparse +import importlib +import sys +import threading +from pathlib import Path + + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument("--entrypoint", required=True) + parser.add_argument("--storage-root", required=True) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--instance-id", required=True) + parser.add_argument("--parent-pid", type=int, required=True) + return parser.parse_args() + + +def load_server(entrypoint): + path = Path(entrypoint).resolve() + if path.name == "__init__.py": + module_root = path.parent.parent + module_name = path.parent.name + else: + module_root = path.parent + module_name = path.stem + addon_root = Path(__file__).resolve().parent.parent + for directory in (module_root, addon_root): + if str(directory) not in sys.path: + sys.path.insert(0, str(directory)) + module = importlib.import_module(module_name) + server = getattr(module, "server", None) + if server is None: + raise RuntimeError(f"Entrypoint does not export 'server': {entrypoint}") + return server + + +def main(): + args = parse_args() + server = load_server(args.entrypoint) + server.bind(args.storage_root) + + import uvicorn + from blendjob.server import watch_parent + + app = server.create_app(args.instance_id) + config = uvicorn.Config( + app, + host=args.host, + port=args.port, + access_log=False, + log_level="warning", + ) + http_server = uvicorn.Server(config) + + def stop_http_server(): + server.shutdown_event.wait() + http_server.should_exit = True + + threading.Thread(target=stop_http_server, daemon=True).start() + threading.Thread( + target=watch_parent, + args=(args.parent_pid, server.shutdown_event), + name="BlendJobParentMonitor", + daemon=True, + ).start() + http_server.run() + + +if __name__ == "__main__": + main() diff --git a/src/snapform/blendjob/runtime.py b/src/snapform/blendjob/runtime.py new file mode 100644 index 0000000..7781840 --- /dev/null +++ b/src/snapform/blendjob/runtime.py @@ -0,0 +1,596 @@ +import json +import os +import subprocess +import sys +import threading +from pathlib import Path + +from .controller import ServerController +from .environment import environment_digest, normalized_environment +from .operator import JobOperatorBase + + +class EnvironmentController: + """Install the declared Environment, then run arbitrary project setup.""" + + def __init__(self, runtime): + self.runtime = runtime + self.process = None + self.log_file = None + self.post_thread = None + self.post_status = {} + self.post_error = "" + self.cancelled = False + self._post_job_id = "" + self._post_stage = 1 + self._lock = threading.Lock() + + def submit(self, _job_type, _parameters): + runtime = self.runtime + runtime.storage_root() + self.post_thread = None + self.post_status = {} + self.post_error = "" + self.cancelled = False + self._post_job_id = "" + self._post_stage = 1 + self.log_file = runtime.environment_log_path().open( + "w", encoding="utf-8" + ) + environment = os.environ.copy() + environment["PYTHONIOENCODING"] = "utf-8:backslashreplace" + environment["PYTHONUTF8"] = "1" + self.process = subprocess.Popen( + runtime.install_command(), + stdout=self.log_file, + stderr=subprocess.STDOUT, + cwd=runtime.storage_root(), + creationflags=windows_creation_flags(), + env=environment, + ) + return { + "job_id": "environment", + "directory": str(runtime.storage_root()), + } + + def status(self, _job_id): + status = self._read_status() + if self.process is None or self.process.poll() is None: + return { + **status, + "state": "running", + "progress": status.get("progress", 0.0), + } + if self.process.returncode != 0: + return { + "message": "Environment installation failed", + **status, + "state": "failed", + "error": status.get("error", "Environment installer failed"), + } + if self.runtime.post_install is None: + return self._success_status() + if self.post_thread is None: + self.post_thread = threading.Thread( + target=self._run_post_install, + name="BlendJobPostInstall", + daemon=True, + ) + self.post_thread.start() + if self.post_thread.is_alive(): + with self._lock: + post_status = dict(self.post_status) + return { + "progress": post_status.get("progress", 0.0), + "message": post_status.get( + "message", + "Running post-install setup", + ), + "stage": post_status.get("stage", 2), + "stage_label": post_status.get("stage_label", "Post Install"), + "state": "running", + } + return self._success_status() + + def cancel(self, _job_id): + self.cancelled = True + if self.process is not None and self.process.poll() is None: + self.process.terminate() + if self._post_job_id: + try: + self.runtime.server.cancel(self._post_job_id) + except RuntimeError: + pass + return {"state": "cancelled"} + + def mark_job_complete(self, _job_id): + if self.log_file is not None: + self.log_file.close() + self.log_file = None + + def _run_post_install(self): + runtime = self.runtime + runtime._request_progress = self._update_post_progress + runtime._request_cancel_check = lambda: self.cancelled + try: + runtime.post_install(runtime) + except Exception as error: + self.post_error = str(error) + finally: + runtime._request_progress = None + runtime._request_cancel_check = None + + def _update_post_progress(self, status): + job_id = str(status.get("job_id", "")) + with self._lock: + if job_id and job_id != self._post_job_id: + self._post_job_id = job_id + self._post_stage += 1 + self.post_status = { + **status, + "stage": self._post_stage, + "stage_label": str(status.get("job_type", "Post Install")), + } + + def _success_status(self): + if self.cancelled: + return { + "progress": 1.0, + "message": "Environment installation cancelled", + "state": "cancelled", + } + if self.post_error: + return { + "progress": 1.0, + "message": f"Environment is ready; post-install failed: {self.post_error}", + "state": "succeeded", + "post_install_error": self.post_error, + } + return { + "progress": 1.0, + "message": "Environment is ready", + "state": "succeeded", + } + + def _read_status(self): + try: + return json.loads( + self.runtime.environment_status_path().read_text( + encoding="utf-8" + ) + ) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return {} + + +class BlenderJobRuntime: + """Own one add-on's Storage Root, Environment, Server and Blender UI.""" + + def __init__( + self, + server, + *, + storage_root, + environment, + post_install=None, + namespace=None, + ): + self.server_definition = server + self.storage_root_factory = storage_root + self.environment = normalized_environment(environment) + self.environment_hash = environment_digest(self.environment) + if post_install is not None and not callable(post_install): + raise TypeError("post_install must be callable") + self.post_install = post_install + self.namespace = namespace or _default_namespace(server.name) + self.active_job = None + self.progress = 0.0 + self.message = "Ready" + self._operator_base = None + self._operator_classes = None + self._timer_registered = False + self._registered = False + self._status_bar_draw = self._create_status_bar_draw() + self._request_progress = None + self._request_cancel_check = None + self._environment_controller = EnvironmentController(self) + self.server = ServerController( + self._server_command, + cwd_factory=self.storage_root, + log_path_factory=self.server_log_path, + creationflags=windows_creation_flags(), + ) + self.JobOperatorBase = self.operator_base() + + def storage_root(self): + configured = ( + self.storage_root_factory() + if callable(self.storage_root_factory) + else self.storage_root_factory + ) + root = Path(configured).expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + return root + + def environment_directory(self): + return self.storage_root() / ".venv" + + def environment_python(self): + directory = self.environment_directory() + if os.name == "nt": + return directory / "Scripts" / "python.exe" + return directory / "bin" / "python" + + def environment_manifest_path(self): + return self.storage_root() / "manifest.json" + + def environment_config_path(self): + return self.storage_root() / "environment.json" + + def environment_status_path(self): + return self.storage_root() / "environment-status.json" + + def environment_log_path(self): + return self.storage_root() / "environment.log" + + def server_log_path(self): + return self.storage_root() / "server.log" + + def environment_manifest(self): + try: + value = json.loads( + self.environment_manifest_path().read_text(encoding="utf-8") + ) + except (FileNotFoundError, json.JSONDecodeError, OSError): + return {} + return value if isinstance(value, dict) else {} + + def environment_ready(self): + return ( + self.environment_python().is_file() + and self.environment_manifest().get("environment_hash") + == self.environment_hash + ) + + def write_environment_config(self): + path = self.environment_config_path() + path.write_text( + json.dumps(self.environment, indent=2), + encoding="utf-8", + ) + return path + + def install_command(self): + installer = Path(__file__).with_name("installer") / "environment.py" + config = self.write_environment_config() + return [ + sys.executable, + str(installer), + "--storage-root", + str(self.storage_root()), + "--config", + str(config), + "--status", + str(self.environment_status_path()), + "--stages", + "1", + ] + + def request(self, job_type, parameters, response=None): + """Directly request one Server Job and return its JobResult.""" + if self._request_cancel_check is not None and self._request_cancel_check(): + raise RuntimeError("Server request was cancelled") + return self.server.request( + job_type, + parameters, + response=response, + progress=self._request_progress, + cancel_check=self._request_cancel_check, + ) + + def operator_base(self): + if self._operator_base is not None: + return self._operator_base + runtime = self + + class BoundJobOperator(JobOperatorBase): + job_runtime = runtime + + BoundJobOperator.__name__ = "JobOperatorBase" + self._operator_base = BoundJobOperator + return BoundJobOperator + + def operator_classes(self): + if self._operator_classes is not None: + return self._operator_classes + import bpy + + runtime = self + operator_base = self.JobOperatorBase + + class InstallEnvironment(operator_base, bpy.types.Operator): + bl_idname = f"{runtime.namespace}.install_environment" + bl_label = "Install Environment" + job_type = "install-environment" + starting_message = "Starting Environment installation" + + def request(self, _context): + if not bpy.app.online_access: + raise RuntimeError( + "Online Access is disabled in Blender Preferences" + ) + return {} + + def controller(self, _runtime): + runtime.server.stop() + return runtime._environment_controller + + class CancelJob(bpy.types.Operator): + bl_idname = f"{runtime.namespace}.cancel_job" + bl_label = "Cancel Job" + + @classmethod + def poll(cls, _context): + return runtime.active_job is not None + + def execute(self, _context): + runtime.cancel_active() + return {"FINISHED"} + + class StartServer(bpy.types.Operator): + bl_idname = f"{runtime.namespace}.start_server" + bl_label = "Start Server" + + def execute(self, _context): + try: + runtime.server.enable() + except (OSError, RuntimeError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + class StopServer(bpy.types.Operator): + bl_idname = f"{runtime.namespace}.stop_server" + bl_label = "Stop Server" + + def execute(self, _context): + runtime.server.stop() + runtime.redraw_ui() + return {"FINISHED"} + + class RestartServer(bpy.types.Operator): + bl_idname = f"{runtime.namespace}.restart_server" + bl_label = "Restart Server" + + def execute(self, _context): + try: + runtime.server.restart() + except (OSError, RuntimeError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} + + class OpenServerLog(bpy.types.Operator): + bl_idname = f"{runtime.namespace}.open_server_log" + bl_label = "Open Server Log" + + def execute(self, _context): + path = runtime.server_log_path() + target = path if path.is_file() else path.parent + result = bpy.ops.wm.path_open(filepath=str(target)) + if "FINISHED" not in result: + self.report({"ERROR"}, "Unable to open the Server log") + return {"CANCELLED"} + return {"FINISHED"} + + self._operator_classes = ( + InstallEnvironment, + CancelJob, + StartServer, + StopServer, + RestartServer, + OpenServerLog, + ) + return self._operator_classes + + def register(self): + if self._registered: + return + import bpy + + for class_type in self.operator_classes(): + bpy.utils.register_class(class_type) + bpy.types.STATUSBAR_HT_header.append(self._status_bar_draw) + self._registered = True + self.enable() + + def unregister(self): + if not self._registered: + return + import bpy + + self.disable() + bpy.types.STATUSBAR_HT_header.remove(self._status_bar_draw) + for class_type in reversed(self.operator_classes()): + bpy.utils.unregister_class(class_type) + self._registered = False + + def begin_job(self, job): + if self.active_job is not None: + raise RuntimeError("Another task is running") + self.active_job = job + + def finish_job(self, job): + if job.job_id: + job.controller.mark_job_complete(job.job_id) + try: + job.operator.cleanup() + finally: + if self.active_job is job: + self.active_job = None + + def cancel_active(self): + job = self.active_job + if job is None: + return + job.cancelled = True + if not job.started or not job.job_id: + return + try: + job.controller.cancel(job.job_id) + except RuntimeError: + pass + + def close_active(self): + job = self.active_job + if job is not None: + self.finish_job(job) + + def update_ui(self, context, progress, message): + self.progress = min(max(float(progress), 0.0), 1.0) + self.message = str(message) + window_manager = getattr(context, "window_manager", None) + if window_manager is not None and hasattr( + window_manager, "progress_update" + ): + window_manager.progress_update(int(self.progress * 100)) + self.redraw_ui(context) + + def redraw_ui(self, context=None, force=False): + try: + import bpy + except ModuleNotFoundError: + return + context = context or bpy.context + window_manager = getattr(context, "window_manager", None) + for window in getattr(window_manager, "windows", ()): + screen = getattr(window, "screen", None) + for area in getattr(screen, "areas", ()): + if area.type in {"STATUSBAR", "VIEW_3D", "PREFERENCES"}: + area.tag_redraw() + for region in getattr(area, "regions", ()): + region.tag_redraw() + if not force: + return + workspace = getattr(context, "workspace", None) + if workspace is not None: + try: + workspace.status_text_set_internal(None) + except (AttributeError, RuntimeError, TypeError): + pass + + def server_status(self): + return dict(self.server.snapshot) + + def server_busy(self): + return self.active_job is not None or self.server.snapshot.get( + "state" + ) == "BUSY" + + def resource(self, name): + return self.server.resource(name) + + def clear_resource(self, name): + return self.server.clear_resource(name) + + def enable(self): + try: + import bpy + except ModuleNotFoundError: + return + timers = getattr(bpy.app, "timers", None) + if timers is None or timers.is_registered(self._poll_server): + return + timers.register(self._poll_server, first_interval=1.0, persistent=True) + self._timer_registered = True + + def disable(self): + self.cancel_active() + self.close_active() + try: + import bpy + except ModuleNotFoundError: + bpy = None + timers = getattr(getattr(bpy, "app", None), "timers", None) + if timers is not None and timers.is_registered(self._poll_server): + timers.unregister(self._poll_server) + self._timer_registered = False + self.server.detach() + + def _poll_server(self): + interval = self.server.poll(available=self.environment_ready()) + self.redraw_ui() + return interval + + def _draw_status_bar(self, owner, _context): + if self.active_job is None: + return + row = owner.layout.row(align=True) + progress_row = row.row(align=True) + progress_row.ui_units_x = 18.0 + progress_row.progress( + factor=self.progress, + type="BAR", + text=self.message, + ) + row.operator( + f"{self.namespace}.cancel_job", + text="", + icon="X", + ) + + def _create_status_bar_draw(self): + runtime = self + + def draw_status_bar(owner, context): + runtime._draw_status_bar(owner, context) + + return draw_status_bar + + def _server_command(self, port, instance_id): + entrypoint_path = self._server_entrypoint() + if entrypoint_path is None: + raise RuntimeError("Job Server does not export a Server entrypoint") + runner = Path(__file__).with_name("runner.py") + return [ + str(self.environment_python()), + "-u", + str(runner), + "--entrypoint", + str(entrypoint_path), + "--storage-root", + str(self.storage_root()), + "--host", + "127.0.0.1", + "--port", + str(port), + "--instance-id", + instance_id, + "--parent-pid", + str(os.getpid()), + ] + + def _server_entrypoint(self): + candidates = [] + for module in tuple(sys.modules.values()): + if getattr(module, "server", None) is not self.server_definition: + continue + module_path = getattr(module, "__file__", None) + if module_path: + candidates.append(Path(module_path).resolve()) + package_entries = [path for path in candidates if path.name == "__init__.py"] + if package_entries: + return min(package_entries, key=lambda path: len(path.parts)) + if candidates: + return min(candidates, key=lambda path: len(path.parts)) + return self.server_definition.entrypoint_path + + +def _default_namespace(server_name): + return str(server_name).split()[0].lower().replace("_", "-") + + +def windows_creation_flags(): + if os.name != "nt": + return 0 + return subprocess.CREATE_NO_WINDOW diff --git a/src/snapform/blendjob/server.py b/src/snapform/blendjob/server.py index b7d8bb3..e886632 100644 --- a/src/snapform/blendjob/server.py +++ b/src/snapform/blendjob/server.py @@ -1,8 +1,10 @@ import os +import sys import threading import time import traceback import uuid +from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from pathlib import Path @@ -17,10 +19,19 @@ class JobCancelled(RuntimeError): class JobContext: """Mutable state and cancellation API passed to a registered job handler.""" - def __init__(self, job_id, job_type, parameters, resources, directory): + def __init__( + self, + job_id, + job_type, + parameters, + resources, + storage_root, + directory, + ): self.job_id = job_id self.job_type = job_type self.parameters = parameters + self.storage_root = Path(storage_root) self.directory = Path(directory) self._resources = resources self._lock = threading.Lock() @@ -34,8 +45,7 @@ class JobContext: "directory": str(self.directory), } - def update(self, progress, message, error=None, state=None, **details): - """Publish a progress snapshot returned by ``GET /jobs/{job_id}``.""" + def _update(self, progress, message, error=None, state=None, **details): if state is None: if error: state = "failed" @@ -57,6 +67,15 @@ class JobContext: with self._lock: self._status = snapshot + def progress(self, progress, message, **details): + """Publish running progress from a Job handler.""" + self._update(progress, message, state="running", **details) + + def succeed(self, result=None, message="Task complete"): + """Publish a successful terminal result.""" + details = {"result": result} if result is not None else {} + self._update(1.0, message, state="succeeded", **details) + def snapshot(self): with self._lock: return dict(self._status) @@ -86,26 +105,59 @@ class JobContext: try: return self._resources[name] except KeyError: - raise KeyError(f"Unknown service resource: {name}") from None + raise KeyError(f"Unknown Server Resource: {name}") from None class JobServer: """Small FastAPI wrapper with decorator-based job registration.""" - def __init__(self, name, *, jobs_directory=None, max_job_history=100): + def __init__( + self, + name, + *, + storage_root=None, + max_job_history=100, + ): self.name = name self.max_job_history = max(int(max_job_history), 0) self.handlers = {} self.resources = {} + self.resource_factories = {} self.jobs = {} + self.futures = {} self.active_job_id = None self.lock = threading.Lock() + self.executor = ThreadPoolExecutor( + max_workers=1, + thread_name_prefix="BlendJobWorker", + ) self.shutdown_event = threading.Event() self.started_at = time.time() - self.jobs_directory = Path(jobs_directory or Path.cwd() / "jobs") + self.storage_root = Path(storage_root) if storage_root is not None else None + self.jobs_directory = ( + self.storage_root / "jobs" if self.storage_root is not None else None + ) + self.entrypoint_module = None + self.entrypoint_path = None + + def bind(self, storage_root): + """Bind the Server to the Runtime-owned Storage Root.""" + storage_root = Path(storage_root).resolve() + if self.storage_root is not None: + current = self.storage_root.resolve() + if current != storage_root: + raise RuntimeError( + f"Job Server is already bound to another Storage Root: {current}" + ) + self.storage_root = storage_root + self.jobs_directory = storage_root / "jobs" + for name, factory in tuple(self.resource_factories.items()): + if name not in self.resources: + self.resources[name] = factory(self) + return self def add_resource(self, name, resource): - """Register shared state owned for the full service lifetime.""" + """Register shared state owned for the full Server lifetime.""" if not name or "/" in name: raise ValueError(f"Invalid resource name: {name}") with self.lock: @@ -114,6 +166,20 @@ class JobServer: self.resources[name] = resource return resource + def resource(self, name): + """Register a Resource factory initialized after Storage Root binding.""" + def register(factory): + if not name or "/" in name: + raise ValueError(f"Invalid resource name: {name}") + if name in self.resources or name in self.resource_factories: + raise ValueError(f"Resource is already registered: {name}") + self.resource_factories[name] = factory + if self.storage_root is not None: + self.resources[name] = factory(self) + return factory + + return register + def resource_snapshot(self, name): try: resource = self.resources[name] @@ -129,9 +195,9 @@ class JobServer: } def clear_resource(self, name): - """Clear an idle service resource without racing a new job submission.""" + """Clear an idle Server Resource without racing a new Job submission.""" with self.lock: - if self.active_job_id is not None: + if self.active_job_id is not None or self._queued_job_count(): return False try: resource = self.resources[name] @@ -144,6 +210,7 @@ class JobServer: return True def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) for resource in reversed(tuple(self.resources.values())): close = getattr(resource, "close", None) if close is not None: @@ -155,6 +222,12 @@ class JobServer: if job_type in self.handlers: raise ValueError(f"Job type is already registered: {job_type}") self.handlers[job_type] = handler + if self.entrypoint_module is None: + self.entrypoint_module = handler.__module__ + module = sys.modules.get(handler.__module__) + module_path = getattr(module, "__file__", None) + if module_path: + self.entrypoint_path = Path(module_path).resolve() return handler return register @@ -162,26 +235,28 @@ class JobServer: def snapshot(self, instance_id): with self.lock: active = self.jobs.get(self.active_job_id) + queued_jobs = self._queued_job_count() result = { "ready": True, - "service": self.name, + "server": self.name, "instance_id": instance_id, - "busy": active is not None, + "busy": active is not None or queued_jobs > 0, "active_job": active.snapshot() if active else None, + "queued_jobs": queued_jobs, "started_at": self.started_at, "resources": self.resource_snapshots(), } return result def submit(self, job_type, parameters, *, job_id=None): + if self.storage_root is None or self.jobs_directory is None: + raise RuntimeError("Bind the Job Server to a Storage Root before submitting") job_id = job_id or uuid.uuid4().hex if job_type not in self.handlers: raise KeyError(job_type) if not job_id or Path(job_id).name != job_id or job_id in {".", ".."}: raise ValueError(f"Invalid job id: {job_id}") with self.lock: - if self.active_job_id is not None: - return None if job_id in self.jobs: raise ValueError(f"Job id already exists: {job_id}") self.jobs_directory.mkdir(parents=True, exist_ok=True) @@ -195,13 +270,33 @@ class JobServer: job_type, parameters, self.resources, + self.storage_root, directory, ) self.jobs[job_id] = context - self.active_job_id = job_id self._prune_jobs() + self.futures[job_id] = self.executor.submit(self.run, context) return context + def cancel(self, job_id): + with self.lock: + context = self.jobs.get(job_id) + future = self.futures.get(job_id) + if context is None: + return None + if future is not None and future.cancel(): + context.request_cancel() + context._update(1.0, "Task cancelled", state="cancelled") + self.futures.pop(job_id, None) + return context + return context if context.request_cancel() else None + + def _queued_job_count(self): + return sum( + context.snapshot()["state"] == "queued" + for context in self.jobs.values() + ) + def _prune_jobs(self): completed = [ identifier @@ -214,25 +309,42 @@ class JobServer: self.jobs.pop(identifier, None) def run(self, context): - context.update(0.0, "Service accepted the task", state="running") + with self.lock: + if context.snapshot()["state"] in TERMINAL_STATES: + return + self.active_job_id = context.job_id + context._update(0.0, "Server accepted the task", state="running") try: context.check_cancelled() - self.handlers[context.job_type](context, context.parameters) + result = self.handlers[context.job_type]( + context, + context.parameters, + ) context.check_cancelled() - if context.snapshot()["state"] not in TERMINAL_STATES: - context.update(1.0, "Task complete", state="succeeded") + snapshot = context.snapshot() + if snapshot["state"] not in {"failed", "cancelled"}: + context.succeed( + result, + message=snapshot.get("message") or "Task complete", + ) except JobCancelled: - context.update(1.0, "Task cancelled", state="cancelled") + context._update(1.0, "Task cancelled", state="cancelled") except Exception as error: traceback.print_exc() - context.update(1.0, "Task failed", error=error, state="failed") + context._update(1.0, "Task failed", error=error, state="failed") finally: with self.lock: if self.active_job_id == context.job_id: self.active_job_id = None + self.futures.pop(context.job_id, None) - def create_app(self, instance_id): - from fastapi import BackgroundTasks, FastAPI, HTTPException + def create_app(self, instance_id, storage_root=None): + from fastapi import FastAPI, HTTPException + + if storage_root is not None: + self.bind(storage_root) + elif self.storage_root is None: + raise RuntimeError("Storage Root is required to create the Job Server app") @asynccontextmanager async def lifespan(_app): @@ -276,14 +388,11 @@ class JobServer: except TypeError as error: raise HTTPException(status_code=405, detail=str(error)) if not cleared: - raise HTTPException(status_code=409, detail="Service is busy") + raise HTTPException(status_code=409, detail="Server is busy") return {"resource": resource_name, "state": "cleared"} @app.post("/jobs", status_code=202) - def submit_job( - payload: dict, - background_tasks: BackgroundTasks, - ): + def submit_job(payload: dict): job_type = str(payload.get("job_type", payload.get("command", ""))) parameters = payload.get("parameters") if not job_type or not isinstance(parameters, dict): @@ -294,9 +403,6 @@ class JobServer: raise HTTPException(status_code=422, detail="Unknown job type") except ValueError as error: raise HTTPException(status_code=409, detail=str(error)) - if context is None: - raise HTTPException(status_code=409, detail="Service is busy") - background_tasks.add_task(self.run, context) return context.snapshot() @app.get("/jobs/{job_id}") @@ -308,8 +414,8 @@ class JobServer: @app.delete("/jobs/{job_id}", status_code=202) def cancel_job(job_id: str): - context = self.jobs.get(job_id) - if context is None or not context.request_cancel(): + context = self.cancel(job_id) + if context is None: raise HTTPException(status_code=404, detail="Active job was not found") return context.snapshot() @@ -322,7 +428,7 @@ class JobServer: def watch_parent(parent_pid, shutdown_event, interval=1.0): - """Stop a dedicated service after its owning Blender process exits.""" + """Stop a dedicated Server after its owning Blender process exits.""" if not parent_pid: return while not shutdown_event.wait(interval): diff --git a/src/snapform/client.py b/src/snapform/client.py deleted file mode 100644 index ef47d67..0000000 --- a/src/snapform/client.py +++ /dev/null @@ -1,521 +0,0 @@ -import json -import os -import platform -import subprocess -import sys -import tempfile -import threading -from pathlib import Path - -import bpy - -from .blendjob import JobRequest, ServiceController -from .env_revision import MANIFEST_FILENAME, MANIFEST_REVISION -from .models import ( - BEN2_MODEL, - MOGE2_MODEL, - get_downloadable_model, - model_directory_name, - selected_model_key, -) - - -SERVICE_ADDRESS = "127.0.0.1" -DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" -_service_controller = None -_environment_controller = None - - -class EnvironmentController: - """Adapt the local Environment bootstrap to the modal controller protocol.""" - - def __init__(self): - self.process = None - self.log_file = None - - def submit(self, _request): - stop_service() - paths = install_paths() - self.log_file = paths["log"].open("w", encoding="utf-8") - environment = os.environ.copy() - environment["PYTHONIOENCODING"] = "utf-8:backslashreplace" - environment["PYTHONUTF8"] = "1" - self.process = subprocess.Popen( - install_env_command(paths["status"]), - stdout=self.log_file, - stderr=subprocess.STDOUT, - cwd=storage_root(), - creationflags=creation_flags(), - env=environment, - ) - return {"job_id": "environment", "directory": str(storage_root())} - - def status(self, _job_id): - status = read_bootstrap_status(install_paths()["status"]) or {} - if self.process is None or self.process.poll() is None: - return {**status, "state": "running", "progress": status.get("progress", 0.0)} - if self.process.returncode == 0: - return {"progress": 1.0, "message": "Environment is ready", **status, - "state": "succeeded"} - return {"message": "Environment installation failed", **status, - "state": "failed", "error": status.get("error", "Bootstrap failed")} - - def cancel(self, _job_id): - if self.process is not None and self.process.poll() is None: - self.process.terminate() - return {"state": "cancelled"} - - def mark_job_complete(self, _job_id): - if self.log_file is not None: - self.log_file.close() - self.log_file = None - - -def environment_controller(): - global _environment_controller - if _environment_controller is None: - _environment_controller = EnvironmentController() - return _environment_controller - - -def preferences(): - addon = bpy.context.preferences.addons.get(__package__) - return addon.preferences if addon is not None else None - - -def storage_root(): - addon_preferences = preferences() - configured = getattr(addon_preferences, "storage_root", "") - root = ( - Path(bpy.path.abspath(configured)).expanduser() - if configured - else DEFAULT_STORAGE_ROOT - ) - root.mkdir(parents=True, exist_ok=True) - return root - - -def venv_directory(): - return storage_root() / ".venv" - - -def venv_python(): - directory = venv_directory() - if os.name == "nt": - return directory / "Scripts" / "python.exe" - return directory / "bin" / "python" - - -def env_installed(): - ready_path = storage_root() / MANIFEST_FILENAME - if not venv_python().is_file() or not ready_path.is_file(): - return False - try: - manifest = json.loads(ready_path.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - return False - return ( - isinstance(manifest, dict) - and manifest.get("manifest_revision") == MANIFEST_REVISION - ) - - -def models_directory(): - path = storage_root() / "models" - path.mkdir(parents=True, exist_ok=True) - return path - - -def model_directory(model_key): - root = models_directory() - directory = root / model_directory_name(model_key) - old_moge_directory = root / "moge2_base_normal" - if ( - model_key == "MOGE2_VITB_NORMAL" - and not directory.is_dir() - and old_moge_directory.is_dir() - ): - old_moge_directory.rename(directory) - legacy_directory = root / model_key.lower() - if ( - not directory.is_dir() - and legacy_directory.is_dir() - and legacy_directory != directory - ): - legacy_directory.rename(directory) - return directory - - -def model_ready(model_key): - model = get_downloadable_model(model_key) - path = model_directory(model_key) - return all( - any(path.glob(pattern)) - for pattern in model.ready_patterns - ) - - -def env_devices(): - """Devices detected by the installed environment, or None when unknown.""" - ready_path = storage_root() / MANIFEST_FILENAME - if not ready_path.is_file(): - return None - try: - payload = json.loads(ready_path.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - return None - devices = payload.get("devices") - if not isinstance(devices, list) or not devices: - return None - return [str(name).lower() for name in devices] - - -def machine_devices(): - """Devices the machine can run before environment installation.""" - devices = ["cpu"] - if sys.platform == "win32": - devices.insert(0, "directml") - elif sys.platform == "darwin" and platform.machine().lower() in { - "arm64", - "aarch64", - }: - devices.insert(0, "mps") - return devices - - -def supported_devices(): - """Inference devices available on this machine, fastest first.""" - return env_devices() or machine_devices() - - -DEVICE_LABELS = { - "cuda": ("CUDA", "NVIDIA CUDA"), - "directml": ("DirectML", "Windows GPU (DirectML)"), - "mps": ("MPS", "Apple Metal (MPS)"), - "cpu": ("CPU", "CPU"), -} - - -def device_enum_items(_owner, _context): - """Blender EnumProperty items listing only devices supported on this machine.""" - return [ - (name.upper(), *DEVICE_LABELS[name], index) - for index, name in enumerate(supported_devices()) - ] - - -def production_device(): - """Device configured for production tasks, falling back to Auto.""" - addon_preferences = preferences() - configured = getattr(addon_preferences, "device", "") - return configured or "AUTO" - - -def temporary_input_path(): - directory = Path(tempfile.mkdtemp(prefix="snapform-input-")) - return directory / "input.png" - - -def temporary_input_directory(): - return Path(tempfile.mkdtemp(prefix="snapform-input-")) - - -def install_paths(): - """Return shared bootstrap paths outside the BlendJob workspace.""" - root = storage_root() - return { - "job_id": "", - "log": root / "environment.log", - "status": root / "environment-status.json", - "result": root, - } - - -def uv_executable(): - """Return the downloaded uv path inside the shared storage.""" - return str(uv_cache_path()) - - -def uv_cache_path(): - if os.name == "nt": - return storage_root() / "tools" / "uv.exe" - return storage_root() / "tools" / "uv" - - -def install_env_command(status_path): - executable = uv_executable() - script = Path(__file__).with_name("install_env.py") - command = [ - sys.executable, - str(script), - "--uv", - executable, - "--venv-dir", - str(venv_directory()), - "--manifest", - str(storage_root() / MANIFEST_FILENAME), - "--models-dir", - str(models_directory()), - "--status", - str(status_path), - ] - return command - - -def service_log_path(): - return storage_root() / "service.log" - - -def _service_command(port, instance_id): - return [ - str(venv_python()), - "-u", - str(Path(__file__).with_name("ai_service") / "server.py"), - "--host", - SERVICE_ADDRESS, - "--port", - str(port), - "--instance-id", - instance_id, - "--parent-pid", - str(os.getpid()), - ] - - -def service_controller(): - global _service_controller - if _service_controller is None: - _service_controller = ServiceController( - _service_command, - cwd_factory=storage_root, - log_path_factory=service_log_path, - host=SERVICE_ADDRESS, - creationflags=creation_flags(), - ) - return _service_controller - - -def start_service(): - return service_controller().start() - - -def ensure_service(timeout=60.0): - connection = service_controller().ensure(timeout) - _redraw_service_ui() - return connection - - -def stop_service(): - service_controller().stop() - - -def restart_service(): - return service_controller().restart() - - -def enable_service(): - return service_controller().enable() - - -def clear_service_models(): - return ensure_service().client.clear_resource("models") - - -def detach_service_client(): - service_controller().detach() - - -def schedule_service_start(): - timers = getattr(bpy.app, "timers", None) - if timers is None: - return - if timers.is_registered(_service_timer): - return - timers.register(_service_timer, first_interval=1.0, persistent=True) - - -def unschedule_service(): - timers = getattr(bpy.app, "timers", None) - if timers is None: - return - if timers.is_registered(_service_timer): - timers.unregister(_service_timer) - - -def service_status(): - return dict(service_controller().snapshot) - - -def service_busy(): - return service_controller().snapshot.get("state") == "BUSY" - - -def _service_timer(): - interval = service_controller().poll(available=env_installed()) - _redraw_service_ui() - return interval - - -def _redraw_service_ui(): - if threading.current_thread() is not threading.main_thread(): - return - try: - from .status_bar import redraw_ui - - redraw_ui(bpy.context, force=True) - except (AttributeError, RuntimeError): - pass - - -def download_model(model_key): - get_downloadable_model(model_key) - return JobRequest("download-model", {"model": model_key}) - - -def remove_background(settings, input_path, model_key, *, start_frame=1, device=None): - return _background_request( - "remove-background", settings, input_path, model_key, start_frame, device - ) - - -def debug_ben2(settings, input_path, model_key, *, start_frame=1, device=None): - return _background_request( - "debug-ben2", settings, input_path, model_key, start_frame, device - ) - - -def _background_request(job_type, settings, input_path, model_key, start_frame, device): - return JobRequest(job_type, { - "input": str(input_path), - "model": model_key, - "device": (device or settings.device).lower(), - "video_frames": getattr(settings, "video_frames", 1), - "start_frame": start_frame, - }) - - -def extract_frames(settings, input_path, *, start_frame=1): - return JobRequest("extract-frames", { - "input": str(input_path), - "video_frames": settings.video_frames, - "start_frame": start_frame, - }) - - -def generate_da3_depth(settings, input_path, model_key, *, process_res, use_ray_pose, - remove_background=False, start_frame=1, device=None): - return JobRequest("generate-da3-depth", { - **_da3_common(settings, input_path, model_key, process_res, device, start_frame), - "export_format": "mini_npz", - "use_ray_pose": 1 if use_ray_pose else 0, - "inference_mode": "per_frame", - "remove_background_model": BEN2_MODEL.key if remove_background else "", - }) - - -def generate_da3_surface(settings, input_path, model_key, *, process_res, use_ray_pose, - remove_background=False, start_frame=1, device=None): - return JobRequest("generate-da3-surface", { - **_da3_common(settings, input_path, model_key, process_res, device, start_frame), - "export_format": "mesh_data", - "use_ray_pose": 1 if use_ray_pose else 0, - "inference_mode": "per_frame", - "remove_background_model": BEN2_MODEL.key if remove_background else "", - "mesh_filter_edges": bool(getattr(settings, "mesh_filter_edges", True)), - "mesh_max_edge_length": float(getattr(settings, "mesh_max_edge_length", 0.1)), - }) - - -def debug_da3(settings, input_path=None, model_key=None, *, device=None): - export_format = settings.export_format - selected_key = model_key or selected_model_key(settings, export_format) - return JobRequest("debug-da3", { - **_da3_common( - settings, - input_path or Path(bpy.path.abspath(settings.input_path)), - selected_key, - settings.process_res, - device, - 1, - ), - "export_format": export_format, - "use_ray_pose": 1 if settings.use_ray_pose else 0, - "inference_mode": "joint", - "remove_background_model": "", - "mesh_filter_edges": bool(getattr(settings, "mesh_filter_edges", True)), - "mesh_max_edge_length": float(getattr(settings, "mesh_max_edge_length", 0.1)), - }) - - -def _da3_common(settings, input_path, model_key, process_res, device, start_frame): - return { - "input": str(input_path), - "model": model_key, - "process_res": process_res, - "process_res_method": "upper_bound_resize", - "conf_thresh_percentile": settings.conf_thresh_percentile, - "num_max_points": settings.num_max_points, - "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": (device or settings.device).lower(), - "ref_view_strategy": settings.ref_view_strategy, - } - - -def generate_moge2_depth(settings, input_path, model_key=MOGE2_MODEL.key, *, - resolution_level=5, remove_background=False, - start_frame=1, device=None): - return JobRequest("generate-moge2-depth", { - **_moge2_common(settings, input_path, model_key, resolution_level, device, start_frame), - "remove_background_model": BEN2_MODEL.key if remove_background else "", - }) - - -def generate_moge2_surface(settings, input_path, model_key=MOGE2_MODEL.key, *, - resolution_level=5, remove_background=False, - start_frame=1, device=None): - return JobRequest("generate-moge2-surface", { - **_moge2_common(settings, input_path, model_key, resolution_level, device, start_frame), - "remove_background_model": BEN2_MODEL.key if remove_background else "", - }) - - -def debug_moge2(settings, input_path=None, model_key=None, *, device=None): - return JobRequest("debug-moge2", { - **_moge2_common( - settings, - input_path or Path(bpy.path.abspath(settings.input_path)), - model_key or settings.moge_model, - settings.moge_resolution_level, - device, - 1, - ), - "export_format": settings.moge_export_format, - "include_points": True, - }) - - -def _moge2_common(settings, input_path, model_key, resolution_level, device, start_frame): - return { - "input": str(input_path), - "model": model_key, - "resolution_level": int(resolution_level), - "video_frames": settings.video_frames, - "start_frame": start_frame, - "device": (device or settings.device).lower(), - } - - -def creation_flags(): - if os.name != "nt": - return 0 - return subprocess.CREATE_NO_WINDOW - - -def read_bootstrap_status(path): - try: - return json.loads(Path(path).read_text(encoding="utf-8")) - except (FileNotFoundError, json.JSONDecodeError, OSError): - return None diff --git a/src/snapform/env_revision.py b/src/snapform/env_revision.py deleted file mode 100644 index d0a6201..0000000 --- a/src/snapform/env_revision.py +++ /dev/null @@ -1,4 +0,0 @@ -"""Shared environment metadata for the Blender UI and the install subprocess.""" - -MANIFEST_FILENAME = "manifest.json" -MANIFEST_REVISION = 8 diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py deleted file mode 100644 index daf077c..0000000 --- a/src/snapform/install_env.py +++ /dev/null @@ -1,921 +0,0 @@ -import argparse -import hashlib -import http.client -import json -import os -import platform -import queue -import shutil -import subprocess -import sys -import tarfile -import tempfile -import threading -import time -import urllib.request -import zipfile -from pathlib import Path - - -SCRIPT_DIRECTORY = Path(__file__).resolve().parent -if str(SCRIPT_DIRECTORY) not in sys.path: - sys.path.insert(0, str(SCRIPT_DIRECTORY)) - -from env_revision import MANIFEST_REVISION -from models import BEN2_MODEL, DOWNLOADABLE_MODELS, MOGE2_MODEL - - -PYPI_MIRRORS = ( - "https://pypi.tuna.tsinghua.edu.cn/simple", - "https://mirrors.aliyun.com/pypi/simple/", -) -PYPI_DEFAULT_INDEX = "https://pypi.org/simple/" -FASTAPI_VERSION = "0.139.2" -UVICORN_VERSION = "0.51.0" -ONNXRUNTIME_VERSION = "1.24.4" -NUMPY_VERSION = "2.4.2" -PILLOW_VERSION = "12.1.1" -AV_VERSION = "17.0.0" -TRIMESH_VERSION = "4.11.2" -HUGGINGFACE_HUB_VERSION = "1.4.1" -UV_VERSION = "0.11.3" -UV_OFFICIAL_CONNECT_TIMEOUT = 3.0 -UV_MIRROR_CONNECT_TIMEOUT = 15.0 -DOWNLOAD_CHUNK_SIZE = 256 * 1024 -UV_OFFICIAL_BASE_URL = ( - f"https://github.com/astral-sh/uv/releases/download/{UV_VERSION}" -) -UV_MIRROR_BASE_URL = f"https://mirror.omoolab.xyz/uv/{UV_VERSION}" -UV_ARTIFACTS = { - "x86_64-pc-windows-msvc": ( - "uv-x86_64-pc-windows-msvc.zip", - "ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c", - ), - "aarch64-apple-darwin": ( - "uv-aarch64-apple-darwin.tar.gz", - "2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f", - ), - "x86_64-unknown-linux-gnu": ( - "uv-x86_64-unknown-linux-gnu.tar.gz", - "c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce", - ), -} -_resolved_pypi_index = None -DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key) -RAMP_RATE = 0.02 -CREEP_CAP = 0.60 -CHECK_CODE = ( - "import json; import fastapi, uvicorn, onnxruntime; " - "import av, numpy, PIL, trimesh; " - "print('SNAPFORM_DEVICES=' + json.dumps({" - "'onnxruntime': onnxruntime.__version__, " - "'providers': onnxruntime.get_available_providers()}))" -) - - -class _ConnectHTTPSConnection(http.client.HTTPSConnection): - """Limit TCP/TLS setup without limiting artifact reads.""" - - connect_timeout = UV_OFFICIAL_CONNECT_TIMEOUT - - def connect(self): - try: - self.timeout = self.connect_timeout - super().connect() - finally: - self.timeout = None - self.sock.settimeout(None) - - -class _MirrorConnectHTTPSConnection(_ConnectHTTPSConnection): - connect_timeout = UV_MIRROR_CONNECT_TIMEOUT - - -class _ConnectHTTPSHandler(urllib.request.HTTPSHandler): - def __init__(self, connection_type): - super().__init__() - self.connection_type = connection_type - - def https_open(self, request): - return self.do_open(self.connection_type, request) - - -def build_uv_download_opener(connection_type): - return urllib.request.build_opener( - urllib.request.ProxyHandler(), - _ConnectHTTPSHandler(connection_type), - ) - - -_uv_download_openers = { - UV_OFFICIAL_BASE_URL: build_uv_download_opener(_ConnectHTTPSConnection), - UV_MIRROR_BASE_URL: build_uv_download_opener( - _MirrorConnectHTTPSConnection - ), -} - - -def configure_standard_streams(): - """Windows 图形进程可能继承 GBK stdout，安装日志统一写为 UTF-8。""" - for stream in (sys.stdout, sys.stderr): - reconfigure = getattr(stream, "reconfigure", None) - if callable(reconfigure): - reconfigure(encoding="utf-8", errors="backslashreplace") - - -def write_status( - path, - progress, - message, - error=None, - stage=None, - stages=None, - stage_label=None, -): - payload = {"progress": progress, "message": message} - if error: - payload["error"] = error - if stage is not None: - payload["stage"] = stage - payload["stages"] = stages - payload["stage_label"] = stage_label - Path(path).write_text( - json.dumps(payload, ensure_ascii=False), - encoding="utf-8", - ) - - -def child_environment(environment=None): - """Inherit the current environment and inject system/registry proxy settings into child processes (uv only reads proxy env variables).""" - process_environment = os.environ.copy() - process_environment["PYTHONIOENCODING"] = "utf-8:backslashreplace" - process_environment["PYTHONUTF8"] = "1" - for key, value in urllib.request.getproxies().items(): - env_key = f"{key.upper()}_PROXY" - if env_key not in process_environment: - process_environment[env_key] = value - if _resolved_pypi_index: - process_environment["UV_INDEX_URL"] = _resolved_pypi_index - if environment: - process_environment.update(environment) - return process_environment - - -def run(command, environment=None): - print("Running:", subprocess.list2cmdline([str(value) for value in command])) - subprocess.check_call( - [str(value) for value in command], - env=child_environment(environment), - ) - - -def parse_package_count(line): - try: - return int(line.split()[1]) - except (IndexError, ValueError): - return 0 - - -def resolve_indexes(args): - """确定用户指定的 PyPI 源，优先命令行参数，其次 uv/pip 的环境变量。""" - global _resolved_pypi_index - _resolved_pypi_index = ( - getattr(args, "pypi_index", None) - or os.environ.get("UV_INDEX_URL") - or os.environ.get("PIP_INDEX_URL") - ) - - -def _run_install_once( - status_path, - stage, - stages, - label, - command, - environment, - start, - end, - clock, -): - """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" - process = subprocess.Popen( - [str(value) for value in command], - env=process_environment, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - encoding="utf-8", - errors="replace", - ) - output_queue = queue.Queue() - - def read_output(): - for raw_line in process.stdout: - output_queue.put(raw_line) - output_queue.put(None) - - threading.Thread(target=read_output, daemon=True).start() - - def eased(fraction): - return start + min(max(fraction, 0.0), 1.0) * (end - start) - - installed = [] - total = 0 - displayed = start - target = start - started = clock() - last_tick = started - last_heartbeat = 0.0 - def detail(message): - write_status( - status_path, - displayed, - message, - stage=stage, - stages=stages, - stage_label=label, - ) - - while True: - try: - raw_line = output_queue.get(timeout=0.2) - empty = False - except queue.Empty: - raw_line = None - empty = True - now = clock() - elapsed = now - started - dt = now - last_tick - last_tick = now - if total: - effective = max(target, eased(CREEP_CAP)) - if effective > displayed: - displayed = min(effective, displayed + dt * RAMP_RATE) - if raw_line is None and not empty: - break - if empty: - if elapsed - last_heartbeat >= 5.0: - last_heartbeat = elapsed - # uv 的下载阶段没有增量进度，只报告已解析的包总数， - # 让用户对规模有预期，而不承诺剩余时间或已耗时。 - message = ( - f"Downloading {total} packages" - if total - else "Downloading packages" - ) - write_status( - status_path, - displayed, - message, - stage=stage, - stages=stages, - stage_label=label, - ) - continue - line = raw_line.rstrip() - print(line) - if line.startswith("Resolved ") and "package" in line: - total = parse_package_count(line) - resolved_progress = eased(0.02) - displayed = max(displayed, resolved_progress) - target = max(target, resolved_progress) - detail(f"Resolving {total} packages") - elif line.startswith("Prepared ") and "package" in line: - done = parse_package_count(line) - if total: - target = max(target, eased(0.10 + 0.60 * done / total)) - detail(f"Downloading packages {done}/{total}") - elif line.startswith("Installed ") and "package" in line: - done = parse_package_count(line) - if total: - target = max(target, eased(0.75 + 0.25 * done / total)) - detail(f"Installing packages {done}/{total}") - elif line.startswith("+"): - entry = line[1:].strip() - installed.append(entry) - return_code = process.wait() - if return_code != 0: - raise subprocess.CalledProcessError(return_code, command) - 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 across PyPI indexes on failure.""" - if _resolved_pypi_index is None: - # 未指定源时优先国内镜像，pypi.org 作为最终兜底。 - indexes = (*PYPI_MIRRORS, PYPI_DEFAULT_INDEX) - elif _resolved_pypi_index == PYPI_DEFAULT_INDEX: - # 明确选择官方源时，失败后回退国内镜像。 - indexes = (PYPI_DEFAULT_INDEX, *PYPI_MIRRORS) - else: - # 用户指定了特定源，不再回退。 - indexes = (_resolved_pypi_index,) - attempts = [ - dict(environment or {}, UV_INDEX_URL=index) for index in indexes - ] - for position, 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 position + 1 == len(attempts): - raise - print( - "Retrying with next PyPI index: " - f"{attempts[position + 1]['UV_INDEX_URL']}" - ) - - -def venv_python(venv_dir): - if sys.platform == "win32": - return Path(venv_dir) / "Scripts" / "python.exe" - return Path(venv_dir) / "bin" / "python" - - -def install_stages(args, venv_dir, python, uv): - """Build the lightweight ONNX-only runtime and model download stages.""" - runtime_package = ( - f"onnxruntime-directml=={ONNXRUNTIME_VERSION}" - if sys.platform == "win32" - else f"onnxruntime=={ONNXRUNTIME_VERSION}" - ) - stages = [ - ( - "Environment", - [uv, "venv", "--clear", "--python", "3.12", venv_dir], - None, - "plain", - ), - ( - "ONNX Runtime", - [ - uv, - "pip", - "install", - "-v", - "--python", - python, - f"fastapi=={FASTAPI_VERSION}", - f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}", - f"numpy=={NUMPY_VERSION}", - f"av=={AV_VERSION}", - runtime_package, - f"pillow=={PILLOW_VERSION}", - f"trimesh=={TRIMESH_VERSION}", - f"uvicorn=={UVICORN_VERSION}", - ], - None, - "install", - ), - ] - models = DOWNLOADABLE_MODELS - worker_script = Path(__file__).with_name("ai_service") / "jobs.py" - models_dir = Path(args.models_dir) - for model_key in DEFAULT_MODEL_KEYS: - model = models[model_key] - download_command = [ - python, - worker_script, - "download", - "--repository", - model.repository, - "--model-dir", - models_dir / model.directory_name, - "--status", - args.status, - "--label", - model.label, - ] - for pattern in model.download_patterns: - download_command.extend(["--allow-pattern", pattern]) - for pattern in model.ready_patterns: - download_command.extend(["--required-pattern", pattern]) - if model.r2_directory: - download_command.extend( - ["--r2-directory", model.r2_directory] - ) - for filename, size, checksum in model.r2_files: - download_command.extend( - ["--r2-file", f"{filename}|{size}|{checksum}"] - ) - stages.append( - ( - model.label, - download_command, - None, - "download", - ) - ) - return stages - - -def env_installed(venv_dir, manifest_path=None): - ready_path = Path(manifest_path or Path(venv_dir).parent / "manifest.json") - if not venv_python(venv_dir).is_file() or not ready_path.is_file(): - return False - try: - payload = json.loads(ready_path.read_text(encoding="utf-8")) - except (json.JSONDecodeError, OSError): - return False - return ( - isinstance(payload, dict) - and payload.get("manifest_revision") == MANIFEST_REVISION - ) - - -def installed_devices(): - devices = [] - machine = platform.machine().lower() - if sys.platform == "darwin" and machine in {"arm64", "aarch64"}: - devices.append("mps") - elif sys.platform == "win32": - devices.append("directml") - devices.append("cpu") - return devices - - -def verify_command(python): - return [python, "-c", CHECK_CODE] - - -def run_verify( - python, - status_path, - start=0.0, - end=1.0, - clock=time.monotonic, - stage=None, - stages=None, - stage_label=None, -): - """Run the environment import check with heartbeat progress; return the device report.""" - command = verify_command(python) - print("Running:", subprocess.list2cmdline([str(value) for value in command])) - process = subprocess.Popen( - [str(value) for value in command], - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - encoding="utf-8", - errors="replace", - ) - output_queue = queue.Queue() - - def read_output(): - for raw_line in process.stdout: - output_queue.put(raw_line) - output_queue.put(None) - - threading.Thread(target=read_output, daemon=True).start() - - report = {} - started = clock() - last_heartbeat = 0.0 - while True: - try: - raw_line = output_queue.get(timeout=0.2) - empty = False - except queue.Empty: - raw_line = None - empty = True - if raw_line is None and not empty: - break - if empty: - elapsed = clock() - started - if elapsed - last_heartbeat >= 5.0: - last_heartbeat = elapsed - creep = min(elapsed / 60.0, CREEP_CAP) - progress = start + creep * (end - start) - write_status( - status_path, - progress, - "Checking Environment", - stage=stage, - stages=stages, - stage_label=stage_label, - ) - continue - line = raw_line.rstrip() - print(line) - if line.startswith("SNAPFORM_DEVICES="): - try: - report = json.loads(line[len("SNAPFORM_DEVICES=") :]) - except json.JSONDecodeError: - report = {} - return_code = process.wait() - if return_code != 0: - raise subprocess.CalledProcessError(return_code, command) - return report - - -def uv_download_target(): - machine = platform.machine().lower() - if sys.platform == "win32": - if machine in {"arm64", "aarch64"}: - return "aarch64-pc-windows-msvc", "zip" - return "x86_64-pc-windows-msvc", "zip" - if sys.platform == "darwin": - arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64" - return f"{arch}-apple-darwin", "tar.gz" - if sys.platform.startswith("linux"): - arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64" - return f"{arch}-unknown-linux-gnu", "tar.gz" - raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}") - - -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): - executable_name = "uv.exe" if sys.platform == "win32" else "uv" - with tempfile.TemporaryDirectory() as directory: - if zipfile.is_zipfile(archive_path): - with zipfile.ZipFile(archive_path) as archive: - archive.extractall(directory) - else: - with tarfile.open(archive_path) as archive: - archive.extractall(directory) - source = next(Path(directory).rglob(executable_name)) - shutil.copy2(source, uv_path) - if sys.platform != "win32": - os.chmod(uv_path, 0o755) - - -def file_sha256(path): - digest = hashlib.sha256() - with Path(path).open("rb") as stream: - while chunk := stream.read(1024 * 1024): - digest.update(chunk) - return digest.hexdigest() - - -def format_download_rate(rate): - size = float(rate) - units = ("B/s", "KB/s", "MB/s", "GB/s") - unit_index = 0 - while size >= 1024 and unit_index < len(units) - 1: - size /= 1024 - unit_index += 1 - precision = 0 if unit_index == 0 else 1 - return f"{size:.{precision}f} {units[unit_index]}" - - -def download_uv_archive( - destination, - status_path, - start, - end, - stage_fields, - clock=time.monotonic, -): - target, _archive_type = uv_download_target() - artifact = UV_ARTIFACTS.get(target) - if artifact is None: - raise RuntimeError(f"SnapForm does not provide uv for {target}") - filename, expected_checksum = artifact - failures = [] - sources = ( - (UV_OFFICIAL_BASE_URL, ""), - (UV_MIRROR_BASE_URL, " (mirror)"), - ) - for base_url, source_label in sources: - destination.unlink(missing_ok=True) - url = f"{base_url}/{filename}" - write_status( - status_path, - start, - f"Downloading uv{source_label}", - **stage_fields, - ) - request = urllib.request.Request( - url, - headers={"User-Agent": "SnapForm uv bootstrap"}, - ) - try: - with _uv_download_openers[base_url].open(request) as response: - total = int(response.headers.get("Content-Length", 0)) - downloaded = 0 - started_at = clock() - last_report_at = started_at - with destination.open("wb") as output: - while chunk := response.read(DOWNLOAD_CHUNK_SIZE): - output.write(chunk) - downloaded += len(chunk) - now = clock() - if now - last_report_at < 0.20 and downloaded < total: - continue - last_report_at = now - elapsed = max(now - started_at, 1e-6) - rate = downloaded / elapsed - message = ( - f"Downloading uv{source_label} " - f"{format_download_rate(rate)}" - ) - progress = start - if total: - progress += (end - start) * downloaded / total - write_status( - status_path, - min(progress, end), - message, - **stage_fields, - ) - actual_checksum = file_sha256(destination) - if actual_checksum != expected_checksum: - raise RuntimeError( - f"SHA-256 mismatch: {actual_checksum}" - ) - if not archive_is_valid(destination): - raise RuntimeError("Downloaded uv archive is invalid") - return destination - except Exception as error: - destination.unlink(missing_ok=True) - failures.append(f"{url}: {error}") - raise RuntimeError("Unable to download uv; " + "; ".join(failures)) - - -def install_uv( - uv, - status_path, - stage=None, - stages=None, - stage_label=None, - start=0.0, - end=1.0, -): - """Download, verify, and extract the pinned uv binary.""" - stage_fields = { - "stage": stage, - "stages": stages, - "stage_label": stage_label, - } - Path(uv).parent.mkdir(parents=True, exist_ok=True) - archive_path = Path(uv).with_name("uv-download.tmp") - try: - download_uv_archive( - archive_path, - status_path, - start, - end, - stage_fields, - ) - write_status( - status_path, - end, - "Extracting uv", - **stage_fields, - ) - extract_uv(archive_path, uv) - finally: - archive_path.unlink(missing_ok=True) - - -def run_stage( - status_path, - stage, - stages, - label, - command, - environment, - kind, - start=0.0, - end=1.0, -): - if stage is None: - opening_message = f"Downloading {label}" - elif kind == "install": - opening_message = "Preparing packages" - else: - opening_message = "" - write_status( - status_path, - start, - opening_message, - stage=stage, - stages=stages, - stage_label=label, - ) - installed = [] - if kind == "install": - installed = run_install( - status_path, - stage, - stages, - label, - command, - environment, - start=start, - end=end, - ) - elif kind == "download": - progress_command = [ - *command, - "--progress-start", - str(start), - "--progress-end", - str(end), - ] - if stage is not None: - progress_command.extend( - [ - "--stage", - str(stage), - "--stages", - str(stages), - "--stage-label", - label, - ] - ) - run(progress_command, environment) - else: - run(command, environment) - write_status( - status_path, - end, - f"{label} ready" if stage is None else "", - stage=stage, - stages=stages, - stage_label=label, - ) - return installed - - -def install(args): - resolve_indexes(args) - venv_dir = Path(args.venv_dir) - manifest_path = Path( - getattr(args, "manifest", venv_dir.parent / "manifest.json") - ) - venv_dir.parent.mkdir(parents=True, exist_ok=True) - if not env_installed(venv_dir, manifest_path): - shutil.rmtree(venv_dir, ignore_errors=True) - manifest_path.unlink(missing_ok=True) - uv = Path(args.uv) - if not uv.is_file(): - resolved = shutil.which(args.uv) - uv = Path(resolved) if resolved else uv - uv_missing = not uv.is_file() - python = venv_python(venv_dir) - stages = install_stages(args, venv_dir, python, uv) - env_stages = [stage for stage in stages if stage[3] != "download"] - model_stages = [stage for stage in stages if stage[3] == "download"] - total_stages = int(uv_missing) + 1 + len(model_stages) - stage = 1 - installed_packages = [] - if uv_missing: - install_uv( - uv, - args.status, - stage=stage, - stages=total_stages, - stage_label="uv", - start=0.0, - end=1.0, - ) - stage += 1 - environment_label = "ONNX Environment" - for index, (_label, command, environment, kind) in enumerate(env_stages): - if index >= 2: - raise RuntimeError("Unexpected ONNX Environment stage") - stage_start, stage_end = ((0.0, 0.05) if index == 0 else (0.05, 1.0)) - stage_installed = run_stage( - args.status, - stage, - total_stages, - environment_label, - command, - environment, - kind, - start=stage_start, - end=stage_end, - ) - if isinstance(stage_installed, list): - installed_packages.extend(stage_installed) - stage += 1 - devices = installed_devices() - venv_dir.mkdir(parents=True, exist_ok=True) - manifest_path.parent.mkdir(parents=True, exist_ok=True) - ready_path = manifest_path - ready_path.write_text( - json.dumps( - { - "manifest_revision": MANIFEST_REVISION, - "python": "3.12", - "onnxruntime": ONNXRUNTIME_VERSION, - "fastapi": FASTAPI_VERSION, - "uvicorn": UVICORN_VERSION, - "runtime_variant": ( - "directml" if sys.platform == "win32" else "default" - ), - "devices": devices, - "packages": sorted(set(installed_packages)), - }, - indent=2, - ), - encoding="utf-8", - ) - failed_models = [] - for label, command, environment, kind in model_stages: - try: - run_stage( - args.status, - stage, - total_stages, - label, - command, - environment, - kind, - start=0.0, - end=1.0, - ) - except Exception as error: - print(f"Model download failed: {label}: {error}") - failed_models.append(label) - write_status( - args.status, - 1.0, - f"Skipping {label}; continuing", - stage=stage, - stages=total_stages, - stage_label=label, - ) - stage += 1 - if failed_models: - labels = ", ".join(failed_models) - write_status( - args.status, - 1.0, - f"Environment is installed; retry {labels} from Models", - stage=total_stages, - stages=total_stages, - stage_label=model_stages[-1][0], - ) - return - write_status( - args.status, - 1.0, - "Environment and recommended models ready", - stage=total_stages, - stages=total_stages, - stage_label=model_stages[-1][0] if model_stages else environment_label, - ) - - -def parse_args(): - parser = argparse.ArgumentParser() - parser.add_argument("--uv", required=True) - parser.add_argument("--venv-dir", required=True) - parser.add_argument("--manifest", required=True) - parser.add_argument("--models-dir", required=True) - parser.add_argument("--status", required=True) - parser.add_argument("--pypi-index", default=None) - return parser.parse_args() - - -def main(): - configure_standard_streams() - args = parse_args() - try: - install(args) - except Exception as error: - write_status( - args.status, - 1.0, - "Environment installation failed", - str(error), - ) - raise - - -if __name__ == "__main__": - main() diff --git a/src/snapform/job_runtime.py b/src/snapform/job_runtime.py new file mode 100644 index 0000000..da5a55c --- /dev/null +++ b/src/snapform/job_runtime.py @@ -0,0 +1,54 @@ +from pathlib import Path + +import bpy + +from .blendjob import BlenderJobRuntime +from .server import server + + +DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" + + +def configured_storage_root(): + addon = bpy.context.preferences.addons.get(__package__) + configured = getattr(getattr(addon, "preferences", None), "storage_root", "") + if configured: + return Path(bpy.path.abspath(configured)) + return DEFAULT_STORAGE_ROOT + + +ENVIRONMENT = { + "python": "3.12", + "packages": [ + "numpy==2.4.2", + "pillow==12.1.1", + "av==17.0.0", + "trimesh==4.11.2", + "huggingface-hub==1.4.1", + ], + "platform_packages": { + "windows": ["onnxruntime-directml==1.24.4"], + "default": ["onnxruntime==1.24.4"], + }, +} + + +def post_install(job_runtime): + errors = [] + for model in ("MOGE2_VITS_NORMAL", "BEN2_BASE"): + try: + job_runtime.request("download-model", {"model": model}) + except RuntimeError as error: + errors.append(str(error)) + if errors: + raise RuntimeError("; ".join(errors)) + + +runtime = BlenderJobRuntime( + server, + storage_root=configured_storage_root, + environment=ENVIRONMENT, + post_install=post_install, +) + +JobOperatorBase = runtime.JobOperatorBase diff --git a/src/snapform/operators.py b/src/snapform/operators.py index acdca04..ee86d7c 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -1,23 +1,14 @@ import shutil -from dataclasses import dataclass from pathlib import Path import bpy -from . import client -from .blendjob import ( - JobRequest, - ModalJobOperatorBase, - active_job, - cancel_active_job, - close_active_job, -) -from .models import ( - BEN2_MODEL, - MOGE2_MODEL, - get_downloadable_model, - get_model, - selected_model_key, +from .job_runtime import JobOperatorBase, runtime +from .properties import ( + BEN2_MODEL_KEY, + DEFAULT_MOGE2_MODEL_KEY, + model_ready, + model_record, ) @@ -35,660 +26,741 @@ BEN2_VIDEO_SUFFIXES = { ".mpg", ".mpeg", } +HIDDEN = {"HIDDEN", "SKIP_SAVE"} -def is_ben2_input(input_path): - return input_path.is_dir() or ( - input_path.suffix.lower() - in BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES +def input_path(value): + path = Path(bpy.path.abspath(value)) + if path.is_file() or path.is_dir(): + return path + raise ValueError("Select a valid image, video, or frame directory") + + +def background_input_path(value): + path = input_path(value) + if path.is_dir() or path.suffix.lower() in ( + BEN2_IMAGE_SUFFIXES | BEN2_VIDEO_SUFFIXES + ): + return path + raise ValueError( + "BEN2 requires an image, video, or directory of image frames" ) -def empty_image_start_frame(source_object): - """Return the first frame to extract, honoring the Empty's Offset within the clip length.""" +def source_start_frame(source_object_name): from .blender.depth_plane import extract_frame_range - return extract_frame_range(source_object)[0] + source_object = bpy.data.objects.get(source_object_name) + return extract_frame_range(source_object)[0] if source_object else 1 -def open_job_directory(result_directory): - result_directory = Path(result_directory) - if not result_directory.is_dir(): - raise RuntimeError(f"Job directory does not exist: {result_directory}") - result = bpy.ops.wm.path_open(filepath=str(result_directory)) +def production_device(): + addon = bpy.context.preferences.addons.get(__package__) + return getattr(getattr(addon, "preferences", None), "device", "AUTO") + + +def require_environment(): + if not runtime.environment_ready(): + raise RuntimeError("Install the SnapForm Environment first") + + +def require_model(model_key): + model = model_record(model_key) + if model is None: + raise RuntimeError(f"Unknown model: {model_key}") + if not model_ready(model_key): + raise RuntimeError(f"Download {model['label']} first") + return model + + +def cleanup_temporary_input(path, delete_input): + if not delete_input or not path: + return + temporary_path = Path(path) + try: + if temporary_path.is_dir(): + shutil.rmtree(temporary_path, ignore_errors=True) + return + temporary_path.unlink(missing_ok=True) + temporary_path.parent.rmdir() + except OSError: + pass + + +def open_job_directory(directory): + directory = Path(directory) + if not directory.is_dir(): + raise RuntimeError(f"Job directory does not exist: {directory}") + result = bpy.ops.wm.path_open(filepath=str(directory)) if "FINISHED" not in result: - raise RuntimeError(f"Unable to open job directory: {result_directory}") + raise RuntimeError(f"Unable to open job directory: {directory}") -@dataclass -class SnapFormJobData: - """SnapForm-only values needed when importing a completed result.""" - - kind: str - export_format: str = "" - source_object_name: str = "" - depth_data_name: str = "" - process_res: int = IMAGE_DEPTH_DEFAULT_PROCESS_RES - delete_source: bool = True - unshaded: bool = False - remove_background: bool = False - mesh_create_camera: bool = False - temporary_input_path: Path | None = None +def push_undo(message): + result = bpy.ops.ed.undo_push(message=message) + if "FINISHED" not in result: + raise RuntimeError(f"Unable to create the {message} Undo step") -class SnapFormJobOperator(ModalJobOperatorBase): - """SnapForm request validation and completed-result integration.""" - - bl_options = {"INTERNAL"} - result_kind = "" - export_format = "" - starting_message = "Starting task" - uses_temporary_input = False - uses_production_device = False - - def __init_subclass__(cls, **kwargs): - super().__init_subclass__(**kwargs) - annotations = {} - for base in reversed(cls.__mro__[1:]): - annotations.update(getattr(base, "__annotations__", {})) - annotations.update(cls.__dict__.get("__annotations__", {})) - if annotations: - cls.__annotations__ = annotations - - def get_service_controller(self): - if self.result_kind == "install_environment": - return client.environment_controller() - return client.service_controller() - - def create_job_request(self, context): - settings = context.scene.snapform_settings - if self.result_kind != "install_environment" and not client.env_installed(): - raise RuntimeError("Install the SnapForm Environment first") - return self.create_request(settings) - - def create_job_data(self, context): - settings = context.scene.snapform_settings - return SnapFormJobData( - kind=self.result_kind, - export_format=self._export_format(settings), - source_object_name=getattr(self, "snapform_source_object_name", ""), - depth_data_name=getattr(self, "snapform_depth_data_name", ""), - process_res=getattr( - self, - "snapform_process_res", - IMAGE_DEPTH_DEFAULT_PROCESS_RES, - ), - delete_source=getattr(self, "snapform_delete_source", True), - unshaded=getattr(self, "snapform_unshaded", False), - remove_background=bool( - getattr(self, "snapform_remove_background", False) - ), - mesh_create_camera=getattr(self, "snapform_mesh_create_camera", False), - temporary_input_path=( - Path(self.snapform_input_path) - if self.uses_temporary_input - and getattr(self, "snapform_delete_input", False) - else None - ), - ) - - def create_request(self, _settings): - raise NotImplementedError - - def _input_path(self, configured_input): - input_path = Path(bpy.path.abspath(configured_input)) - if input_path.is_file() or input_path.is_dir(): - return input_path - raise ValueError("Select a valid image, video, or frame directory") - - def _source_start_frame(self): - source_name = getattr(self, "snapform_source_object_name", "") - source_object = bpy.data.objects.get(source_name) if source_name else None - return empty_image_start_frame(source_object) if source_object else 1 - - def _require_model(self, model_key, *, downloadable=False): - model = ( - get_downloadable_model(model_key) - if downloadable - else get_model(model_key) - ) - if not client.model_ready(model_key): - raise RuntimeError(f"Download {model.label} first") - return model - - def _require_ben2(self): - self._require_model(BEN2_MODEL.key, downloadable=True) - - def _background_input(self, configured_input): - input_path = self._input_path(configured_input) - if not is_ben2_input(input_path): - raise ValueError( - "BEN2 requires an image, video, or directory of image frames" - ) - self._require_ben2() - return input_path - - def _task_device(self, settings): - """Debug panels use the scene device; production tasks use Preferences.""" - if self.uses_production_device: - return client.production_device() - return settings.device - - def _require_online_access(self): - if not bpy.app.online_access: - raise RuntimeError( - "Online Access is disabled in Blender Preferences" - ) - - def _starting_message(self, settings): - return self.starting_message - - def _export_format(self, _settings): - return self.export_format - - def update_job_ui(self, context, progress, message): - settings = context.scene.snapform_settings - settings.progress = progress - settings.status = message - - def redraw_job_ui(self, context, force=False): - from .status_bar import redraw_ui - - redraw_ui(context, force=force) - - def get_starting_message(self, context): - return self._starting_message(context.scene.snapform_settings) - - def cleanup_job_data(self, data): - temporary_path = data.temporary_input_path if data else None - if temporary_path is None: - return - try: - if temporary_path.is_dir(): - shutil.rmtree(temporary_path, ignore_errors=True) - else: - temporary_path.unlink(missing_ok=True) - temporary_path.parent.rmdir() - except OSError: - pass - - def on_job_succeeded(self, context, job, status): - if job.kind in {"debug_da3", "debug_ben2", "debug_moge2"}: - try: - open_job_directory(job.result_directory) - except Exception as open_error: - self.report({"WARNING"}, str(open_error)) - if job.kind == "debug_da3": - message = ( - f"Exported {job.export_format} to {job.result_directory}" - ) - elif job.kind == "debug_ben2": - message = f"Removed background to {job.result_directory}" - else: - message = f"Exported MoGe-2 data to {job.result_directory}" - self.report({"INFO"}, message) - return message - if job.kind == "generate_da3_depth": - self.update_job_ui(context, 0.96, "Building depth plane") - self._build_depth_plane(context, job) - return "Depth plane is ready" - if job.kind == "generate_moge2_depth": - self.update_job_ui(context, 0.96, "Building MoGe-2 depth plane") - self._build_moge2_depth_plane(job) - return "MoGe-2 depth plane is ready" - if job.kind == "generate_moge2_surface": - self.update_job_ui(context, 0.96, "Building MoGe-2 surface") - self._build_moge2_image_mesh(job) - return "MoGe-2 surface is ready" - if job.kind == "generate_da3_surface": - self.update_job_ui(context, 0.96, "Building perspective surface") - self._build_image_mesh(job) - return "Perspective surface is ready" - if job.kind in {"generate_original_background", "extract_original_frames"}: - self.update_job_ui(context, 0.96, "Building depth plane") - self._build_original_depth_plane(job) - return "Depth plane is ready" - if job.kind == "process_image_background": - self.update_job_ui(context, 0.96, "Updating Image Empty") - self._apply_removed_background(job) - return "Image background is removed" - return status.get("message", "Task complete") - - def _build_image_mesh(self, job): - result = bpy.ops.snapform.build_image_mesh( - snapform_source_object_name=job.source_object_name, - snapform_npz_path=str(job.result_directory / "mesh.npz"), - snapform_delete_source=job.delete_source, - snapform_create_camera=job.mesh_create_camera, - snapform_unshaded=job.unshaded, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to build the perspective surface") - undo_result = bpy.ops.ed.undo_push(message="Convert to Surface") - if "FINISHED" not in undo_result: - raise RuntimeError("Unable to create the surface Undo step") - - def _build_moge2_image_mesh(self, job): - result_directory = job.result_directory - mesh_path = result_directory / "mesh.npz" - if not mesh_path.is_file(): - raise RuntimeError("MoGe-2 did not export the official surface mesh") - result = bpy.ops.snapform.build_image_mesh( - snapform_source_object_name=job.source_object_name, - snapform_npz_path=str(mesh_path), - snapform_delete_source=job.delete_source, - snapform_create_camera=job.mesh_create_camera, - snapform_unshaded=job.unshaded, - snapform_moge2=True, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to build the MoGe-2 perspective surface") - undo_result = bpy.ops.ed.undo_push(message="Convert to Surface") - if "FINISHED" not in undo_result: - raise RuntimeError("Unable to create the surface Undo step") - - def _build_depth_plane(self, context, job): - 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=background_path, - snapform_color_frames_dir=( - str(result_directory) if is_sequence else "" - ), - snapform_depth_frames_dir=( - str(result_directory) - if is_sequence - else "" - ), - snapform_delete_source=job.delete_source, - snapform_unshaded=job.unshaded, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to build the depth plane") - undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane") - if "FINISHED" not in undo_result: - 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 == "generate_original_background": - background_path = str( - result_directory - if is_sequence - else result_directory / "foreground.png" - ) - if 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=background_path, - snapform_color_frames_dir=color_frames_dir, - snapform_depth_frames_dir=( - str(result_directory) if is_sequence else "" - ), - snapform_depth_npz_path="", - snapform_delete_source=job.delete_source, - snapform_unshaded=job.unshaded, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to build the depth plane") - undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane") - if "FINISHED" not in undo_result: - raise RuntimeError("Unable to create the depth plane Undo step") - - def _build_moge2_depth_plane(self, job): - result_directory = job.result_directory - is_sequence = bool(list(result_directory.glob("frame.*"))) - foreground_sequence = bool( - list(result_directory.glob("foreground.*.png")) - ) - result = bpy.ops.snapform.build_moge2_depth_plane( - snapform_source_object_name=job.source_object_name, - snapform_npz_path=str(result_directory / "results.npz"), - snapform_background_path=( - str(result_directory) - if job.remove_background and foreground_sequence - else str(result_directory / "foreground.png") - if job.remove_background - else "" - ), - snapform_color_frames_dir=( - str(result_directory) if is_sequence else "" - ), - snapform_data_frames_dir=str(result_directory), - snapform_delete_source=job.delete_source, - snapform_unshaded=job.unshaded, - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to build the MoGe-2 depth plane") - undo_result = bpy.ops.ed.undo_push(message="Convert to Depth Plane") - if "FINISHED" not in undo_result: - raise RuntimeError("Unable to create the depth plane Undo step") - - def _apply_removed_background(self, job, push_undo=True): - 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), - ) - if "FINISHED" not in result: - raise RuntimeError("Unable to update the Image Empty") - if not push_undo: - return - undo_result = bpy.ops.ed.undo_push(message="Remove Background") - if "FINISHED" not in undo_result: - raise RuntimeError("Unable to create the background removal Undo step") - -class ImageInputJobOperator(SnapFormJobOperator): - uses_temporary_input = True - uses_production_device = True - snapform_input_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_delete_input: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, +def build_da3_depth_response( + context, + result, + *, + source_object_name, + depth_data_name, + remove_background, + delete_source, + unshaded, +): + from .blender.depth_plane import ( + build_depth_plane, + create_removed_background_image, + create_sequence_image, + image_base_name, + is_image_empty, ) - def _image_input(self): - return self._input_path(self.snapform_input_path) - - -class GeometryImageJobOperator(ImageInputJobOperator): - snapform_depth_data_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_delete_source: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_unshaded: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_remove_background: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_mesh_create_camera: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - - -class Da3ImageJobOperator(GeometryImageJobOperator): - request_name = "" - snapform_model_key: bpy.props.StringProperty( - default="BASE", - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_use_ray_pose: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_process_res: bpy.props.IntProperty( - default=IMAGE_DEPTH_DEFAULT_PROCESS_RES, - min=1, - max=65_536, - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def create_request(self, settings): - input_path = self._image_input() - model_key = self.snapform_model_key - self._require_model(model_key) - remove_background = bool( - getattr(self, "snapform_remove_background", False) + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + foreground_paths = sorted(directory.glob("foreground.*.png")) + foreground_image = None + if remove_background: + foreground_image = create_removed_background_image( + source_object, + directory if foreground_paths else directory / "foreground.png", + assign_to_source=False, ) - if remove_background: - self._require_ben2() - return getattr(client, self.request_name)( - settings, - input_path, - model_key, - process_res=self.snapform_process_res, - use_ray_pose=self.snapform_use_ray_pose, - remove_background=remove_background, - start_frame=self._source_start_frame(), - device=self._task_device(settings), + color_image = None + if list(directory.glob("frame.*")) and foreground_image is None: + color_image = create_sequence_image( + f"{image_base_name(source_object.data)}.frame", + directory, ) - - -class Moge2ImageJobOperator(GeometryImageJobOperator): - request_name = "" - snapform_model_key: bpy.props.StringProperty( - default="BASE", - options={"HIDDEN", "SKIP_SAVE"}, + build_depth_plane( + context, + source_object, + result.file("depth"), + delete_source, + unshaded=unshaded, + depth_name=depth_data_name, + color_image=color_image, + foreground_image=foreground_image, + depth_frames_directory=( + directory if list(directory.glob("frame.*")) else None + ), ) - snapform_moge_resolution_level: bpy.props.IntProperty( - default=5, - min=0, - max=9, - options={"HIDDEN", "SKIP_SAVE"}, + push_undo("Convert to Depth Plane") + + +def build_moge2_depth_response( + context, + result, + *, + source_object_name, + remove_background, + delete_source, + unshaded, +): + from .blender.depth_plane import ( + build_moge2_depth_plane, + create_removed_background_image, + create_sequence_image, + image_base_name, + is_image_empty, ) - def create_request(self, settings): - input_path = self._image_input() - model_key = self.snapform_model_key or MOGE2_MODEL.key - self._require_model(model_key, downloadable=True) - remove_background = bool( - getattr(self, "snapform_remove_background", False) + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + foreground_paths = sorted(directory.glob("foreground.*.png")) + foreground_image = None + if remove_background: + foreground_image = create_removed_background_image( + source_object, + directory if foreground_paths else directory / "foreground.png", + assign_to_source=False, ) - if remove_background: - self._require_ben2() - return getattr(client, self.request_name)( - settings, - input_path, - model_key, - resolution_level=self.snapform_moge_resolution_level, - remove_background=remove_background, - start_frame=self._source_start_frame(), - device=self._task_device(settings), + color_image = None + if list(directory.glob("frame.*")) and foreground_image is None: + color_image = create_sequence_image( + f"{image_base_name(source_object.data)}.frame", + directory, ) + build_moge2_depth_plane( + context, + source_object, + result.file("depth"), + delete_source, + unshaded=unshaded, + color_image=color_image, + foreground_image=foreground_image, + data_frames_directory=directory, + ) + push_undo("Convert to Depth Plane") -class BackgroundRequestMixin: - request_name = "" +def build_original_depth_response( + context, + result, + *, + source_object_name, + background, + delete_source, + unshaded, +): + from .blender.depth_plane import ( + build_image_depth_plane, + create_removed_background_image, + create_sequence_image, + image_base_name, + is_image_empty, + ) - def background_input(self, settings): - return settings.background_input_path - - def create_request(self, settings): - input_path = self._background_input(self.background_input(settings)) - return getattr(client, self.request_name)( - settings, - input_path, - BEN2_MODEL.key, - start_frame=self._source_start_frame(), - device=self._task_device(settings), + directory = result.directory + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + foreground_image = None + if background: + foreground_paths = sorted(directory.glob("foreground.*.png")) + foreground_image = create_removed_background_image( + source_object, + directory if foreground_paths else directory / "foreground.png", + assign_to_source=False, ) + color_image = None + if list(directory.glob("frame.*")) and foreground_image is None: + color_image = create_sequence_image( + f"{image_base_name(source_object.data)}.frame", + directory, + ) + build_image_depth_plane( + context, + source_object, + delete_source, + unshaded=unshaded, + color_image=color_image, + foreground_image=foreground_image, + depth_frames_directory=( + directory if list(directory.glob("frame.*")) else None + ), + ) + push_undo("Convert to Depth Plane") -class InstallEnvironment(SnapFormJobOperator): - bl_idname = "snapform.install_environment" - bl_label = "Install Environment" - result_kind = "install_environment" - starting_message = "Starting Environment installation" +def build_surface_response( + context, + result, + *, + source_object_name, + delete_source, + create_camera, + unshaded, + moge2, +): + from .blender.depth_plane import is_image_empty + from .blender.surface import build_image_mesh - def create_request(self, _settings): - self._require_online_access() - return JobRequest("install-environment") + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source Image Empty no longer exists") + build_image_mesh( + context, + source_object, + result.file("surface"), + delete_source, + create_camera, + unshaded, + moge2, + ) + push_undo("Convert to Surface") -class DownloadModel(SnapFormJobOperator): +def apply_background_response(result, source_object_name): + from .blender.depth_plane import ( + create_removed_background_image, + is_image_empty, + ) + + source_object = bpy.data.objects.get(source_object_name) + if not is_image_empty(source_object): + raise RuntimeError("The source image empty is no longer available") + foreground_paths = sorted(result.directory.glob("foreground.*.png")) + create_removed_background_image( + source_object, + result.directory if foreground_paths else result.file("foreground"), + ) + push_undo("Remove Background") + + +def da3_parameters(settings, path, model, process_res, use_ray_pose, start_frame): + return { + "input": str(path), + "model": model, + "process_res": process_res, + "process_res_method": "upper_bound_resize", + "conf_thresh_percentile": settings.conf_thresh_percentile, + "num_max_points": settings.num_max_points, + "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": production_device().lower(), + "ref_view_strategy": settings.ref_view_strategy, + "use_ray_pose": 1 if use_ray_pose else 0, + } + + +def moge2_parameters(settings, path, model, resolution_level, start_frame): + return { + "input": str(path), + "model": model, + "resolution_level": int(resolution_level), + "video_frames": settings.video_frames, + "start_frame": start_frame, + "device": production_device().lower(), + } + + +class DownloadModel(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.download_model" bl_label = "Download Model" - result_kind = "download_model" - model_key: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) + job_type = "download-model" - def create_request(self, _settings): - self._require_online_access() - return client.download_model(self.model_key) + model: bpy.props.StringProperty(options=HIDDEN) - def _starting_message(self, _settings): - model = get_downloadable_model(self.model_key) - return f"Starting {model.label} download" + def request(self, _context): + require_environment() + if not bpy.app.online_access: + raise RuntimeError("Online Access is disabled in Blender Preferences") + model = model_record(self.model) + label = model["label"] if model else self.model + self.starting_message = f"Starting {label} download" + return {"model": self.model} -class DebugDa3(SnapFormJobOperator): +class DebugDa3(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_da3" bl_label = "Debug DA3" - result_kind = "debug_da3" + job_type = "debug-da3" - def create_request(self, settings): - input_path = self._input_path(settings.input_path) - model_key = selected_model_key(settings, settings.export_format) - model = self._require_model(model_key) - if settings.export_format == "gs_ply" and not model.supports_gs: - raise RuntimeError( - "PLY export requires a DA3 Giant or Nested Giant model" - ) - return client.debug_da3( - settings, - input_path, - model_key, - device=self._task_device(settings), - ) + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + model = settings.gs_model if settings.export_format == "gs_ply" else settings.model + spec = require_model(model) + if settings.export_format == "gs_ply" and not spec.get("supports_gs"): + raise RuntimeError("PLY export requires a DA3 model with a Gaussian head") + self.export_format = settings.export_format + self.starting_message = f"Starting {self.export_format} export" + return { + **da3_parameters( + settings, + input_path(settings.input_path), + model, + settings.process_res, + settings.use_ray_pose, + 1, + ), + "device": settings.device.lower(), + "export_format": settings.export_format, + "inference_mode": "joint", + "remove_background_model": "", + "mesh_filter_edges": settings.mesh_filter_edges, + "mesh_max_edge_length": settings.mesh_max_edge_length, + } - def _starting_message(self, settings): - return f"Starting {settings.export_format} export" - - def _export_format(self, settings): - return settings.export_format + def response(self, _context, result): + open_job_directory(result.directory) + return f"Exported {self.export_format} to {result.directory}" -class DebugMoge2(SnapFormJobOperator): +class DebugMoge2(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_moge2" bl_label = "Debug MoGe-2" - result_kind = "debug_moge2" - export_format = "moge_npz" + job_type = "debug-moge2" starting_message = "Starting MoGe-2 inference" - def create_request(self, settings): - input_path = self._input_path(settings.input_path) - self._require_model(settings.moge_model, downloadable=True) - return client.debug_moge2( - settings, - input_path, - settings.moge_model, - device=self._task_device(settings), - ) + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + require_model(settings.moge_model) + return { + **moge2_parameters( + settings, + input_path(settings.input_path), + settings.moge_model, + settings.moge_resolution_level, + 1, + ), + "device": settings.device.lower(), + "export_format": settings.moge_export_format, + "include_points": True, + } + + def response(self, _context, result): + open_job_directory(result.directory) + return f"Exported MoGe-2 data to {result.directory}" -class DebugBen2(BackgroundRequestMixin, SnapFormJobOperator): +class DebugBen2(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.debug_ben2" bl_label = "Debug BEN2" - result_kind = "debug_ben2" - request_name = "debug_ben2" + job_type = "debug-ben2" starting_message = "Starting BEN2 background removal" + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + require_model(BEN2_MODEL_KEY) + return { + "input": str(background_input_path(settings.background_input_path)), + "model": BEN2_MODEL_KEY, + "device": settings.device.lower(), + "video_frames": settings.video_frames, + "start_frame": 1, + } -class GenerateDa3Depth(Da3ImageJobOperator): + def response(self, _context, result): + open_job_directory(result.directory) + return f"Removed background to {result.directory}" + + +class GenerateDa3Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_depth" bl_label = "Generate DA3 Depth" - result_kind = "generate_da3_depth" - export_format = "mini_npz" - request_name = "generate_da3_depth" + job_type = "generate-da3-depth" starting_message = "Starting depth estimation" + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + depth_data_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default="BASE", options=HIDDEN) + use_ray_pose: bpy.props.BoolProperty(default=False, options=HIDDEN) + process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) -class GenerateMoge2Depth(Moge2ImageJobOperator): + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **da3_parameters( + settings, + input_path(self.input_path), + self.model, + self.process_res, + self.use_ray_pose, + source_start_frame(self.source_object_name), + ), + "export_format": "mini_npz", + "inference_mode": "per_frame", + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_da3_depth_response( + context, + result, + source_object_name=self.source_object_name, + depth_data_name=self.depth_data_name, + remove_background=self.remove_background, + delete_source=self.delete_source, + unshaded=self.unshaded, + ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateMoge2Depth(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_depth" bl_label = "Generate MoGe-2 Depth" - result_kind = "generate_moge2_depth" - export_format = "moge_npz" - request_name = "generate_moge2_depth" + job_type = "generate-moge2-depth" starting_message = "Starting MoGe-2 inference" + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) + resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) -class GenerateDa3Surface(Da3ImageJobOperator): + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **moge2_parameters( + settings, + input_path(self.input_path), + self.model, + self.resolution_level, + source_start_frame(self.source_object_name), + ), + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + } + + def response(self, context, result): + build_moge2_depth_response( + context, + result, + source_object_name=self.source_object_name, + remove_background=self.remove_background, + delete_source=self.delete_source, + unshaded=self.unshaded, + ) + return "MoGe-2 depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateDa3Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_da3_surface" bl_label = "Generate DA3 Surface" - result_kind = "generate_da3_surface" - export_format = "mesh_data" - request_name = "generate_da3_surface" + job_type = "generate-da3-surface" starting_message = "Starting perspective surface generation" + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default="BASE", options=HIDDEN) + use_ray_pose: bpy.props.BoolProperty(default=False, options=HIDDEN) + process_res: bpy.props.IntProperty(default=1024, min=1, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) -class GenerateMoge2Surface(Moge2ImageJobOperator): + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **da3_parameters( + settings, + input_path(self.input_path), + self.model, + self.process_res, + self.use_ray_pose, + source_start_frame(self.source_object_name), + ), + "export_format": "mesh_data", + "inference_mode": "per_frame", + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + "mesh_filter_edges": settings.mesh_filter_edges, + "mesh_max_edge_length": settings.mesh_max_edge_length, + } + + def response(self, context, result): + build_surface_response( + context, + result, + source_object_name=self.source_object_name, + delete_source=self.delete_source, + create_camera=self.create_camera, + unshaded=self.unshaded, + moge2=False, + ) + return "Perspective surface is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class GenerateMoge2Surface(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_moge2_surface" bl_label = "Generate MoGe-2 Surface" - result_kind = "generate_moge2_surface" - export_format = "moge_npz" - request_name = "generate_moge2_surface" + job_type = "generate-moge2-surface" starting_message = "Starting MoGe-2 inference" + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + model: bpy.props.StringProperty(default=DEFAULT_MOGE2_MODEL_KEY, options=HIDDEN) + resolution_level: bpy.props.IntProperty(default=5, min=0, max=9, options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + create_camera: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + remove_background: bpy.props.BoolProperty(default=False, options=HIDDEN) -class ProcessImageBackground(BackgroundRequestMixin, ImageInputJobOperator): + def request(self, context): + require_environment() + require_model(self.model) + if self.remove_background: + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + **moge2_parameters( + settings, + input_path(self.input_path), + self.model, + self.resolution_level, + source_start_frame(self.source_object_name), + ), + "remove_background_model": ( + BEN2_MODEL_KEY if self.remove_background else "" + ), + "surface_mesh": True, + } + + def response(self, context, result): + build_surface_response( + context, + result, + source_object_name=self.source_object_name, + delete_source=self.delete_source, + create_camera=self.create_camera, + unshaded=self.unshaded, + moge2=True, + ) + return "MoGe-2 surface is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class ProcessImageBackground(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.process_image_background" bl_label = "Process Image Background" - result_kind = "process_image_background" - request_name = "remove_background" + job_type = "remove-background" starting_message = "Starting Image Empty background removal" - uses_temporary_input = True - uses_production_device = True - def background_input(self, _settings): - return self.snapform_input_path + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + "input": str(background_input_path(self.input_path)), + "model": BEN2_MODEL_KEY, + "device": production_device().lower(), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, _context, result): + apply_background_response(result, self.source_object_name) + return "Image background is removed" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) -class GenerateOriginalBackground( - BackgroundRequestMixin, - GeometryImageJobOperator, -): +class GenerateOriginalBackground(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.generate_original_background" bl_label = "Generate Original Background" - result_kind = "generate_original_background" - request_name = "remove_background" + job_type = "remove-background" starting_message = "Starting background removal" - uses_temporary_input = True - uses_production_device = True - def background_input(self, _settings): - return self.snapform_input_path + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + require_model(BEN2_MODEL_KEY) + settings = context.scene.snapform_settings + return { + "input": str(background_input_path(self.input_path)), + "model": BEN2_MODEL_KEY, + "device": production_device().lower(), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, context, result): + build_original_depth_response( + context, + result, + source_object_name=self.source_object_name, + background=True, + delete_source=self.delete_source, + unshaded=self.unshaded, + ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) -class ExtractOriginalFrames(GeometryImageJobOperator): +class ExtractOriginalFrames(JobOperatorBase, bpy.types.Operator): bl_idname = "snapform.extract_original_frames" bl_label = "Extract Original Frames" - result_kind = "extract_original_frames" + job_type = "extract-frames" starting_message = "Extracting source frames" - def create_request(self, settings): - return client.extract_frames( - settings, - self._image_input(), - start_frame=self._source_start_frame(), + input_path: bpy.props.StringProperty(options=HIDDEN) + source_object_name: bpy.props.StringProperty(options=HIDDEN) + delete_source: bpy.props.BoolProperty(default=True, options=HIDDEN) + delete_input: bpy.props.BoolProperty(default=False, options=HIDDEN) + unshaded: bpy.props.BoolProperty(default=False, options=HIDDEN) + + def request(self, context): + require_environment() + settings = context.scene.snapform_settings + return { + "input": str(input_path(self.input_path)), + "video_frames": settings.video_frames, + "start_frame": source_start_frame(self.source_object_name), + } + + def response(self, context, result): + build_original_depth_response( + context, + result, + source_object_name=self.source_object_name, + background=False, + delete_source=self.delete_source, + unshaded=self.unshaded, ) + return "Depth plane is ready" + + def cleanup(self): + cleanup_temporary_input(self.input_path, self.delete_input) + + +class ClearModels(bpy.types.Operator): + bl_idname = "snapform.clear_models" + bl_label = "Clear Models" + + def execute(self, _context): + try: + runtime.clear_resource("model_manager") + except RuntimeError as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + return {"FINISHED"} JOB_OPERATOR_CLASSES = ( - InstallEnvironment, DownloadModel, DebugDa3, DebugMoge2, @@ -700,60 +772,5 @@ JOB_OPERATOR_CLASSES = ( ProcessImageBackground, GenerateOriginalBackground, ExtractOriginalFrames, + ClearModels, ) - - -class CancelJob(bpy.types.Operator): - bl_idname = "snapform.cancel_job" - bl_label = "Cancel SnapForm Task" - - @classmethod - def poll(cls, _context): - return active_job() is not None - - def execute(self, _context): - cancel_active_job() - return {"FINISHED"} - - -class ControlService(bpy.types.Operator): - bl_idname = "snapform.control_service" - bl_label = "Control SnapForm Service" - - action: bpy.props.EnumProperty( - items=( - ("START", "Start", "Start or connect to the Service"), - ("CLEAR", "Clear Models", "Release cached AI models"), - ("RESTART", "Restart", "Restart the Service"), - ("STOP", "Stop", "Stop the Service"), - ("LOG", "Open Log", "Open the Service log"), - ), - options={"HIDDEN", "SKIP_SAVE"}, - ) - - def execute(self, _context): - try: - if self.action == "START": - client.enable_service() - message = "Service is ready" - elif self.action == "CLEAR": - client.clear_service_models() - message = "AI model cache is cleared" - elif self.action == "RESTART": - client.restart_service() - message = "Service is restarting" - elif self.action == "STOP": - client.stop_service() - message = "Service is stopping" - else: - log_path = client.service_log_path() - target = log_path if log_path.is_file() else log_path.parent - result = bpy.ops.wm.path_open(filepath=str(target)) - if "FINISHED" not in result: - raise RuntimeError("Unable to open the Service log") - return {"FINISHED"} - except (OSError, RuntimeError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} - self.report({"INFO"}, message) - return {"FINISHED"} diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 0856814..7da1ceb 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -1,14 +1,16 @@ import bpy -from .models import BEN2_MODEL, MODELS, MOGE2_MODELS, selected_model_key -from .operators import active_job -from . import client -from .client import model_ready, env_installed +from .job_runtime import runtime +from .properties import BEN2_MODEL_KEY, model_catalog, model_ready -class AiServicePanel(bpy.types.Panel): - bl_label = "AI Service" - bl_idname = "SNAPFORM_PT_service" +def selected_model_key(settings): + return settings.gs_model if settings.export_format == "gs_ply" else settings.model + + +class ServerPanel(bpy.types.Panel): + bl_label = "Job Server" + bl_idname = "SNAPFORM_PT_server" bl_space_type = "VIEW_3D" bl_region_type = "UI" bl_category = "SnapForm" @@ -16,8 +18,8 @@ class AiServicePanel(bpy.types.Panel): def draw(self, _context): layout = self.layout - running = active_job() is not None or client.service_busy() - if not env_installed(): + running = runtime.server_busy() + if not runtime.environment_ready(): layout.label( text="Environment is not installed", icon="ERROR", @@ -31,7 +33,7 @@ class AiServicePanel(bpy.types.Panel): ) return - status = client.service_status() + status = runtime.server_status() state = status.get("state", "STOPPED") icons = { "READY": "RADIOBUT_ON", @@ -54,41 +56,37 @@ class AiServicePanel(bpy.types.Panel): if state == "READY": controls = row.row(align=True) operator = controls.operator( - "snapform.control_service", + "snapform.restart_server", text="", icon="FILE_REFRESH", ) - operator.action = "RESTART" operator = controls.operator( - "snapform.control_service", + "snapform.stop_server", text="", icon="PAUSE", ) - operator.action = "STOP" elif state == "STOPPED": - operator = row.operator( - "snapform.control_service", + row.operator( + "snapform.start_server", text="Start", icon="PLAY", ) - operator.action = "START" else: row.label( text=status.get("message", state.title()), icon=icons.get(state, "QUESTION"), ) if state == "ERROR": - operator = row.operator( - "snapform.control_service", + row.operator( + "snapform.start_server", text="Start", icon="PLAY", ) - operator.action = "START" if state in {"READY", "BUSY"}: cached = ( status.get("resources", {}) - .get("models", {}) + .get("model_manager", {}) .get("loaded", {}) ) cached_names = [name for name in cached.values() if name] @@ -102,12 +100,11 @@ class AiServicePanel(bpy.types.Panel): ) clear = cached_row.row(align=True) clear.enabled = state == "READY" - operator = clear.operator( - "snapform.control_service", + clear.operator( + "snapform.clear_models", text="", icon="TRASH", ) - operator.action = "CLEAR" class ModelsPanel(bpy.types.Panel): @@ -121,34 +118,41 @@ class ModelsPanel(bpy.types.Panel): @classmethod def poll(cls, _context): - return env_installed() + return runtime.environment_ready() def draw(self, _context): layout = self.layout - running = active_job() is not None or client.service_busy() + running = runtime.server_busy() - for model in MOGE2_MODELS.values(): - self._draw_model(layout, model, running) - for model in MODELS.values(): - if model.enabled: + models = model_catalog() + for model in models: + if model.get("family") == "moge2": + self._draw_model(layout, model, running) + for model in models: + if model.get("family") == "da3" and model.get("enabled"): self._draw_model(layout, model, running) layout.separator(type="LINE") - self._draw_model(layout, BEN2_MODEL, running) + ben2 = next( + (model for model in models if model.get("key") == BEN2_MODEL_KEY), + None, + ) + if ben2: + self._draw_model(layout, ben2, running) def _draw_model(self, layout, model, running): row = layout.row(align=True) - row.label(text=model.label) - if model_ready(model.key): + row.label(text=model["label"]) + if model.get("ready"): row.label(text="Ready", icon="CHECKMARK") else: download_row = row.row(align=True) - download_row.enabled = env_installed() and not running + download_row.enabled = runtime.environment_ready() and not running operator = download_row.operator( "snapform.download_model", text="Download", icon="IMPORT", ) - operator.model_key = model.key + operator.model = model["key"] class DebugPanel(bpy.types.Panel): @@ -228,8 +232,7 @@ class Da3OutputDebugPanel(bpy.types.Panel): mesh_box.prop(settings, "mesh_max_edge_length") export_row = layout.row() export_row.enabled = ( - active_job() is None - and not client.service_busy() + not runtime.server_busy() and model_ready(selected_model_key(settings)) ) export_row.operator( @@ -258,8 +261,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel): layout.prop(settings, "moge_resolution_level") run_row = layout.row() run_row.enabled = ( - active_job() is None - and not client.service_busy() + not runtime.server_busy() and model_ready(settings.moge_model) ) run_row.operator( @@ -286,9 +288,8 @@ class Ben2RemovalDebugPanel(bpy.types.Panel): layout.prop(settings, "background_input_path") background_row = layout.row() background_row.enabled = ( - active_job() is None - and not client.service_busy() - and model_ready(BEN2_MODEL.key) + not runtime.server_busy() + and model_ready(BEN2_MODEL_KEY) ) background_row.operator( "snapform.debug_ben2", diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 0ee79ec..992e247 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -1,8 +1,7 @@ -from pathlib import Path - import bpy -from .client import device_enum_items +from .job_runtime import DEFAULT_STORAGE_ROOT, runtime +from .properties import device_enum_items class SnapFormPreferences(bpy.types.AddonPreferences): @@ -15,7 +14,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences): "all Blender versions can use the same directory" ), subtype="DIR_PATH", - default=str(Path.home() / ".snapform"), + default=str(DEFAULT_STORAGE_ROOT), ) device: bpy.props.EnumProperty( name="Device", @@ -28,19 +27,16 @@ class SnapFormPreferences(bpy.types.AddonPreferences): ) def draw(self, _context): - from .operators import active_job - from .client import venv_directory, env_installed - layout = self.layout layout.prop(self, "storage_root") layout.prop(self, "device") layout.separator() - layout.label(text=f"Environment: {venv_directory()}") - if env_installed(): + layout.label(text=f"Environment: {runtime.environment_directory()}") + if runtime.environment_ready(): layout.label(text="Environment is installed", icon="CHECKMARK") else: install_row = layout.row() - install_row.enabled = active_job() is None + install_row.enabled = runtime.active_job is None install_row.operator( "snapform.install_environment", text="Install Environment", @@ -48,8 +44,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences): ) layout.separator() operator = layout.operator( - "snapform.control_service", - text="Open Service Log", + "snapform.open_server_log", + text="Open Server Log", icon="TEXT", ) - operator.action = "LOG" diff --git a/src/snapform/properties.py b/src/snapform/properties.py index 080928c..e01bb8b 100644 --- a/src/snapform/properties.py +++ b/src/snapform/properties.py @@ -1,7 +1,142 @@ +import platform +import sys + import bpy -from .models import gs_model_enum_items, model_enum_items, moge2_model_enum_items -from .client import device_enum_items +from .job_runtime import runtime + + +BEN2_MODEL_KEY = "BEN2_BASE" +DEFAULT_DA3_MODEL_KEY = "BASE" +DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL" + + +DEVICE_LABELS = { + "cuda": ("CUDA", "NVIDIA CUDA"), + "directml": ("DirectML", "Windows GPU (DirectML)"), + "mps": ("MPS", "Apple Metal (MPS)"), + "cpu": ("CPU", "CPU"), +} +_models_cache = {} + + +def supported_devices(): + devices = runtime.environment_manifest().get("devices") + if isinstance(devices, list) and devices: + return [str(device).lower() for device in devices] + if sys.platform == "win32": + return ["directml", "cpu"] + if sys.platform == "darwin" and platform.machine().lower() in { + "arm64", + "aarch64", + }: + return ["mps", "cpu"] + return ["cpu"] + + +def device_enum_items(_owner, _context): + return [ + (name.upper(), *DEVICE_LABELS[name], index) + for index, name in enumerate(supported_devices()) + ] + + +def models_status(refresh=False): + global _models_cache + status = runtime.server_status().get("resources", {}).get("model_manager", {}) + if status: + _models_cache = status + if status and not refresh: + return status + if _models_cache and not refresh: + return _models_cache + if not runtime.environment_ready(): + return {} + try: + status = runtime.resource("model_manager") + _models_cache = status + return status + except RuntimeError: + return _models_cache + + +def model_ready(model_key): + return any( + model.get("key") == model_key and model.get("ready") + for model in models_status().get("catalog", ()) + ) + + +def model_catalog(): + return models_status().get("catalog", ()) + + +def model_record(model_key): + return next( + (model for model in model_catalog() if model.get("key") == model_key), + None, + ) + + +def _model_items(models): + return [ + ( + model["key"], + model["label"] if model.get("ready") else f"{model['label']} (Missing)", + f"{model['description']}; license: {model['license']}", + index, + ) + for index, model in enumerate(models) + ] + + +def model_enum_items(_owner, _context): + return _model_items( + model + for model in model_catalog() + if model.get("family") == "da3" and model.get("enabled") + ) + + +def gs_model_enum_items(_owner, _context): + return _model_items( + model + for model in model_catalog() + if model.get("family") == "da3" + and model.get("enabled") + and model.get("supports_gs") + ) + + +def moge2_model_enum_items(_owner, _context): + return _model_items( + model for model in model_catalog() if model.get("family") == "moge2" + ) + + +def depth_plane_model_enum_items(_owner, _context): + models = [ + model + for model in model_catalog() + if model.get("family") in {"moge2", "da3"} and model.get("enabled") + ] + models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY) + items = _model_items(models) + items.append( + ("ORIGINAL", "Original", "Use the original image as depth", len(items)) + ) + return items + + +def surface_model_enum_items(_owner, _context): + models = [ + model + for model in model_catalog() + if model.get("family") in {"moge2", "da3"} + and model.get("enabled") + ] + models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY) + return _model_items(models) class SnapFormSettings(bpy.types.PropertyGroup): @@ -185,16 +320,3 @@ class SnapFormSettings(bpy.types.PropertyGroup): default=1, min=1, ) - status: bpy.props.StringProperty( - name="Status", - default="Ready", - options={"SKIP_SAVE"}, - ) - progress: bpy.props.FloatProperty( - name="Progress", - default=0.0, - min=0.0, - max=1.0, - subtype="FACTOR", - options={"SKIP_SAVE"}, - ) diff --git a/src/snapform/server/__init__.py b/src/snapform/server/__init__.py new file mode 100644 index 0000000..dedfbda --- /dev/null +++ b/src/snapform/server/__init__.py @@ -0,0 +1,64 @@ +"""SnapForm Job Server without Blender dependencies.""" + +import sys +from pathlib import Path + + +SERVER_DIRECTORY = Path(__file__).resolve().parent +ADDON_DIRECTORY = SERVER_DIRECTORY.parent +for import_directory in (SERVER_DIRECTORY, ADDON_DIRECTORY): + if str(import_directory) not in sys.path: + sys.path.insert(0, str(import_directory)) + +try: + from ..blendjob import JobServer +except ImportError: + from blendjob import JobServer + +from .model_manager import ModelManager, ModelSessionCache +from .background import remove_background as _remove_background +from .da3 import generate_da3 as _generate_da3 +from .media import extract_frames as _extract_frames +from .moge2 import generate_moge2 as _generate_moge2 + + +server = JobServer("SnapForm Job Server") + + +@server.resource("model_manager") +def create_model_manager(job_server): + return ModelManager(job_server.storage_root / "models") + + +@server.job("download-model") +def download_model(context, parameters): + manager = context.resource("model_manager") + return manager.download(context, parameters["model"]) + + +@server.job("debug-ben2") +@server.job("remove-background") +def remove_background(context, parameters): + return _remove_background(context, parameters) + + +@server.job("extract-frames") +def extract_frames(context, parameters): + return _extract_frames(context, parameters) + + +@server.job("debug-moge2") +@server.job("generate-moge2-surface") +@server.job("generate-moge2-depth") +def generate_moge2(context, parameters): + return _generate_moge2(context, parameters) + + +@server.job("debug-da3") +@server.job("generate-da3-surface") +@server.job("generate-da3-depth") +def generate_da3(context, parameters): + return _generate_da3(context, parameters) + + +__all__ = ("ModelManager", "ModelSessionCache", "server") diff --git a/src/snapform/server/background.py b/src/snapform/server/background.py new file mode 100644 index 0000000..ad67270 --- /dev/null +++ b/src/snapform/server/background.py @@ -0,0 +1,170 @@ +import shutil +import time +from pathlib import Path + +from .media import ( + IMAGE_SUFFIXES, + check_cancelled, + copy_frames, + prepare_background_input, +) + + +def infer_foreground( + input_path, + model_dir, + device, + model_cache=None, + cancel_check=None, +): + from PIL import Image + + input_path = Path(input_path) + check_cancelled(cancel_check) + if input_path.suffix.lower() not in IMAGE_SUFFIXES: + raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") + from .onnx_ben2 import create_session, infer + + owns_model = model_cache is None + requested_device = getattr(device, "type", device) + if owns_model: + load_started = time.perf_counter() + model = create_session(model_dir, requested_device) + load_ms = (time.perf_counter() - load_started) * 1000.0 + else: + model, load_ms = model_cache.get_ben2(model_dir, device) + + check_cancelled(cancel_check) + with Image.open(input_path) as opened_image: + image = opened_image.convert("RGB") + inference_started = time.perf_counter() + foreground = infer(model, image) + foreground = preserve_source_rgb(image, foreground) + inference_ms = (time.perf_counter() - inference_started) * 1000.0 + check_cancelled(cancel_check) + return foreground, load_ms, inference_ms + + +def preserve_source_rgb(source, foreground): + """Combine the untouched source RGB with BEN2's foreground alpha.""" + source_rgba = source.convert("RGBA") + alpha = foreground.convert("RGBA").getchannel("A") + if alpha.size != source_rgba.size: + from PIL import Image + + alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS) + source_rgba.putalpha(alpha) + return source_rgba + + +def infer_foreground_frames( + images, + model_dir, + device, + output_directory, + progress_callback, + 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) + progress_callback( + 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, + 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(context, parameters): + """Run BEN2 and write transparent foreground images into the Job directory.""" + total_started = time.perf_counter() + context.progress(0.02, "Starting Environment") + context.check_cancelled() + + output_directory = context.directory + output_directory.mkdir(parents=True, exist_ok=True) + device = parameters["device"] + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + + video_frames = parameters.get("video_frames", 16) + start_frame = parameters.get("start_frame", 1) + images, temporary_directory = prepare_background_input( + parameters["input"], + video_frames, + output_directory, + start_frame=start_frame, + cancel_check=context.check_cancelled, + ) + try: + context.progress(0.12, f"Loading BEN2 for {device.upper()}") + if len(images) == 1: + context.progress(0.35, "Removing background") + foreground, load_ms, inference_ms = infer_foreground( + images[0], + model_directory, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + context.progress(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, + model_directory, + device, + output_directory, + context.progress, + 0.12, + 0.7, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + copy_frames(images, output_directory) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + + context.check_cancelled() + 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{frame_text} in {inference_ms:.0f} ms " + f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)" + ) + print(message) + context.progress(1.0, message) + if tuple(output_directory.glob("foreground.*.png")): + return {"foreground": "."} + return {"foreground": "foreground.png"} diff --git a/src/snapform/server/da3.py b/src/snapform/server/da3.py new file mode 100644 index 0000000..d9422be --- /dev/null +++ b/src/snapform/server/da3.py @@ -0,0 +1,233 @@ +import shutil + +from .background import infer_foreground, infer_foreground_frames +from .geometry import ( + export_da3_depth_preview, + export_da3_mini_npz, + export_da3_point_cloud, + export_image_mesh_data, + export_prediction_attributes, + export_textured_mesh_glb, +) +from .media import ( + copy_frames, + flatten_exported_paths, + prepare_input, + wait_for_exported_paths, +) + + +def merge_da3_frame_predictions(predictions): + """Combine independent single-frame DA3 results for existing exporters.""" + import numpy as np + + if not predictions: + raise ValueError("DA3 returned no frame predictions") + merged = predictions[0] + for name in ( + "depth", + "conf", + "extrinsics", + "intrinsics", + "processed_images", + ): + values = [getattr(prediction, name) for prediction in predictions] + setattr(merged, name, np.concatenate(values, axis=0)) + return merged + + +def generate_da3(context, parameters): + """Run DA3 inference and return the artifact used by the requesting Job.""" + context.progress(0.02, "Starting Environment") + context.check_cancelled() + from .onnx_da3 import infer + from .onnx_runtime import provider_label + + output_directory = context.directory + output_directory.mkdir(parents=True, exist_ok=True) + context.progress( 0.08, "Preparing input") + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + device = parameters["device"] + export_format = parameters["export_format"] + images, temporary_directory = prepare_input( + parameters["input"], + parameters.get("video_frames", 16), + output_directory, + start_frame=parameters.get("start_frame", 1), + cancel_check=context.check_cancelled, + ) + + try: + context.check_cancelled() + remove_background_model = parameters.get("remove_background_model", "") + remove_background_model_dir = ( + model_manager.directory(remove_background_model) + if remove_background_model + else "" + ) + foreground_images = [] + if remove_background_model_dir: + context.progress( 0.12, "Removing image background") + if len(images) == 1: + foreground, load_ms, inference_ms = infer_foreground( + images[0], + remove_background_model_dir, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + 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") + foreground_images = [output_directory / "foreground.png"] + else: + load_ms, inference_ms = infer_foreground_frames( + images, + remove_background_model_dir, + device, + output_directory, + context.progress, + 0.12, + 0.05, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + foreground_images = sorted( + output_directory.glob("foreground.*.png") + ) + if remove_background_model_dir: + print( + f"BEN2 inference {inference_ms:.0f} ms; " + f"model load {load_ms:.0f} ms" + ) + context.progress( 0.15, "Loading DA3 ONNX model") + model = model_manager.get_da3(model_directory, device) + context.progress( + 0.20, + f"Loaded DA3 on {provider_label(model)}", + ) + + context.check_cancelled() + context.progress( + 0.25, + f"Estimating depth for {len(images)} frame(s)", + ) + inference_options = { + "process_res": parameters["process_res"], + "process_res_method": parameters["process_res_method"], + "use_ray_pose": bool(parameters["use_ray_pose"]), + "ref_view_strategy": parameters["ref_view_strategy"], + } + try: + if parameters.get("inference_mode", "joint") == "per_frame": + predictions = [] + for index, image in enumerate(images): + context.check_cancelled() + context.progress( + 0.25 + 0.60 * index / len(images), + f"Estimating DA3 depth ({index + 1}/{len(images)})", + ) + predictions.append( + infer(model, [str(image)], **inference_options) + ) + prediction = merge_da3_frame_predictions(predictions) + else: + prediction = infer( + model, + [str(path) for path in images], + **inference_options, + ) + except UnicodeDecodeError as error: + # DirectML reports failures with a locale-specific message that + # ONNX Runtime cannot always decode, masking the real error. + if provider_label(model) == "DirectML": + raise RuntimeError( + "DA3 inference failed on DirectML. Switch the device to " + "CPU in SnapForm Preferences and retry. " + f"Original error: {error}" + ) from error + raise + context.check_cancelled() + context.progress( 0.88, f"Exporting {export_format}") + if export_format == "mesh": + export_textured_mesh_glb( + prediction, + images, + output_directory, + parameters.get("conf_thresh_percentile", 40.0), + bool(parameters.get("mesh_filter_edges", True)), + float(parameters.get("mesh_max_edge_length", 0.1)), + ) + elif export_format == "mesh_data": + def mesh_progress(step, total, label): + context.progress( + 0.88 + 0.09 * step / total, + f"Preparing surface data {step}/{total}: {label}", + ) + + export_image_mesh_data( + prediction, + images[0], + output_directory, + foreground_path=( + foreground_images[0] if foreground_images else None + ), + progress_callback=mesh_progress, + ) + elif export_format == "mini_npz": + export_da3_mini_npz(prediction, output_directory) + elif export_format == "depth_vis": + export_da3_depth_preview(prediction, output_directory) + elif export_format == "glb": + export_da3_point_cloud( + prediction, + output_directory, + parameters["conf_thresh_percentile"], + parameters["num_max_points"], + ) + else: + raise ValueError( + f"DA3 ONNX does not support export format: {export_format}" + ) + if export_format in {"glb", "gs_ply", "depth_vis", "mesh"}: + context.progress( + 0.97, + "Writing structured debug attributes", + ) + export_prediction_attributes(prediction, output_directory) + context.check_cancelled() + + result_paths = wait_for_exported_paths( + output_directory, + export_format, + cancel_check=context.check_cancelled, + ) + if not result_paths: + raise RuntimeError( + f"No {export_format} files were exported" + ) + result_paths = flatten_exported_paths( + output_directory, + result_paths, + cancel_check=context.check_cancelled, + ) + context.check_cancelled() + copy_frames(images, output_directory) + context.progress( + 1.0, + f"Exported {len(result_paths)} {export_format} file(s)", + ) + if context.job_type == "generate-da3-surface": + return {"surface": "mesh.npz"} + if context.job_type == "generate-da3-depth": + return {"depth": "results.npz"} + return {"directory": "."} + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) diff --git a/src/snapform/server/geometry.py b/src/snapform/server/geometry.py new file mode 100644 index 0000000..b0fdbbb --- /dev/null +++ b/src/snapform/server/geometry.py @@ -0,0 +1,616 @@ +from pathlib import Path + + +SCENE_FILENAME = "scene.glb" + + +def unproject_depth_frame(depth, intrinsics, extrinsics): + import numpy as np + + height, width = depth.shape + x, y = np.meshgrid( + np.arange(width, dtype=np.float32), + np.arange(height, dtype=np.float32), + ) + pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) + rays = (np.linalg.inv(intrinsics) @ pixels.T).T + camera_points = rays * depth.reshape(-1, 1) + camera_points = np.column_stack( + (camera_points, np.ones(camera_points.shape[0], dtype=np.float32)) + ) + world_to_camera = np.eye(4, dtype=np.float32) + world_to_camera[: extrinsics.shape[0], :4] = extrinsics + world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T + return world_points[:, :3].reshape(height, width, 3) + + +def unproject_camera_depth(depth, intrinsics): + import numpy as np + + height, width = depth.shape + x, y = np.meshgrid( + np.arange(width, dtype=np.float32), + np.arange(height, dtype=np.float32), + ) + pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3) + rays = (np.linalg.inv(intrinsics) @ pixels.T).T + return (rays * depth.reshape(-1, 1)).reshape(height, width, 3) + + +def compact_mesh(vertices, faces): + import numpy as np + + vertices = np.asarray(vertices, dtype=np.float32) + faces = np.asarray(faces, dtype=np.int32) + used_vertices = np.unique(faces) + index_map = np.full(len(vertices), -1, dtype=np.int32) + index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32) + vertices = vertices[used_vertices] + faces = index_map[faces] + return vertices, faces + + +def project_camera_uv(vertices, intrinsics, image_size): + import numpy as np + + width, height = image_size + z = vertices[:, 2] + if np.any(~np.isfinite(z)) or np.any(z <= 0.0): + raise ValueError("Simplified mesh contains invalid camera depth") + u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2] + v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2] + denominator_x = max(width - 1, 1) + denominator_y = max(height - 1, 1) + return np.column_stack( + ( + np.clip(u / denominator_x, 0.0, 1.0), + np.clip(1.0 - v / denominator_y, 0.0, 1.0), + ) + ).astype(np.float32) + + +def project_camera_pixels(vertices, intrinsics): + import numpy as np + + z = vertices[:, 2] + if np.any(~np.isfinite(z)) or np.any(z <= 0.0): + raise ValueError("Mesh contains invalid camera depth") + return np.column_stack( + ( + intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2], + intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2], + ) + ) + + +def sample_image_nearest(values, pixels): + import numpy as np + + height, width = values.shape + x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1) + y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1) + return np.asarray(values)[y, x] + + +def confidence_percentile_map(depth, confidence): + import numpy as np + + valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) + result = np.zeros(depth.shape, dtype=np.float32) + if not valid.any(): + return result + sorted_values = np.sort(np.asarray(confidence)[valid]) + ranks = np.searchsorted( + sorted_values, + np.asarray(confidence)[valid], + side="right", + ) + result[valid] = ranks.astype(np.float32) / len(sorted_values) + return result + + +def depth_discontinuity_ratio(depth): + import numpy as np + + padded = np.pad(depth, 1, mode="edge") + windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3)) + difference = np.nanmax(windows, axis=(-2, -1)) + difference -= np.nanmin(windows, axis=(-2, -1)) + return difference / np.maximum(depth, 1e-6) + + +def triangle_edge_stretch_values(points, triangles): + import numpy as np + + flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3) + triangle_points = flat_points[triangles] + triangle_depth = triangle_points[..., 2] + safe_depth = np.where( + np.abs(triangle_depth) > 1e-8, + triangle_depth, + np.nan, + ) + triangle_rays = triangle_points / safe_depth[..., None] + next_points = np.roll(triangle_points, -1, axis=1) + next_rays = np.roll(triangle_rays, -1, axis=1) + next_depth = np.roll(triangle_depth, -1, axis=1) + local_depth = (triangle_depth + next_depth) * 0.5 + flat_edges = ( + triangle_rays * local_depth[..., None] + - next_rays * local_depth[..., None] + ) + flat_lengths = np.linalg.norm(flat_edges, axis=-1) + actual_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + stretch = actual_lengths / np.maximum(flat_lengths, 1e-8) + stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0 + return stretch + + +def image_mesh_attributes( + vertices, + faces, + intrinsics, + confidence_percentiles, + face_depth_discontinuity, +): + import numpy as np + + point_percentile = sample_image_nearest( + confidence_percentiles, + project_camera_pixels(vertices, intrinsics), + ).astype(np.float32) + face_confidence = point_percentile[faces].min(axis=1) + return { + "point_depth": vertices[:, 2].astype(np.float32), + "face_confidence_percentile": face_confidence.astype(np.float32), + "face_depth_discontinuity": np.asarray( + face_depth_discontinuity, + dtype=np.float32, + ), + } + + +def export_prediction_attributes(prediction, output_directory): + import numpy as np + + values = {} + for name in ("depth", "conf", "intrinsics", "extrinsics"): + value = getattr(prediction, name, None) + if value is not None: + values[name] = np.asarray(value) + if "depth" in values and "conf" in values: + values["confidence_percentile"] = np.stack( + [ + confidence_percentile_map(depth, confidence) + for depth, confidence in zip( + values["depth"], + values["conf"], + ) + ] + ) + output_path = Path(output_directory) / "prediction.attributes.npz" + np.savez(output_path, **values) + return output_path + + +def export_da3_mini_npz(prediction, output_directory): + import numpy as np + + values = {"depth": np.round(prediction.depth, 8)} + for name in ("conf", "extrinsics", "intrinsics"): + value = getattr(prediction, name, None) + if value is not None: + values[name] = np.round(value, 2) if name == "conf" else value + output_path = Path(output_directory) / "results.npz" + np.savez_compressed(output_path, **values) + return output_path + + +def export_da3_depth_preview(prediction, output_directory): + import numpy as np + from PIL import Image + + target = Path(output_directory) / "depth_vis" + target.mkdir(parents=True, exist_ok=True) + for index, depth in enumerate(prediction.depth): + valid = np.isfinite(depth) & (depth > 0.0) + normalized = np.zeros(depth.shape, dtype=np.float32) + if valid.any(): + low, high = np.percentile(depth[valid], (2.0, 98.0)) + normalized[valid] = np.clip( + (depth[valid] - low) / max(float(high - low), 1e-6), + 0.0, + 1.0, + ) + color = np.stack( + ( + normalized, + 1.0 - np.abs(normalized * 2.0 - 1.0), + 1.0 - normalized, + ), + axis=-1, + ) + color = np.asarray(color * 255.0, dtype=np.uint8) + source = np.asarray(prediction.processed_images[index], dtype=np.uint8) + combined = np.concatenate((source, color), axis=1) + Image.fromarray(combined, mode="RGB").save( + target / f"{index:04d}.jpg", + quality=95, + ) + + +def export_da3_point_cloud( + prediction, + output_directory, + confidence_percentile=40.0, + num_max_points=1_000_000, +): + import numpy as np + import trimesh + + depth = np.asarray(prediction.depth, dtype=np.float32) + confidence = np.asarray(prediction.conf, dtype=np.float32) + threshold = confidence_percentile_threshold( + depth, + confidence, + confidence_percentile, + ) + point_sets = [] + color_sets = [] + for index in range(depth.shape[0]): + points = unproject_depth_frame( + depth[index], + prediction.intrinsics[index], + prediction.extrinsics[index], + ) + valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0) + if threshold is not None: + valid &= confidence[index] >= threshold + point_sets.append(points[valid]) + color_sets.append(prediction.processed_images[index][valid]) + points = np.concatenate(point_sets, axis=0) + colors = np.concatenate(color_sets, axis=0) + maximum = max(int(num_max_points), 1) + if len(points) > maximum: + indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64) + points = points[indices] + colors = colors[indices] + scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors)) + scene.export(Path(output_directory) / SCENE_FILENAME) + + +def remove_background_faces(faces, foreground_mask): + """Return only triangles whose vertices all belong to the foreground.""" + foreground_points = foreground_mask.reshape(-1) + return faces[foreground_points[faces].all(axis=1)] + + +def export_image_mesh_data( + prediction, + image_path, + output_directory, + foreground_path=None, + progress_callback=None, +): + import numpy as np + from PIL import Image + + total_steps = 6 + + def progress(step, label): + if progress_callback is not None: + progress_callback(step, total_steps, label) + + progress(1, "Validating DA3 data") + required = ("depth", "conf", "intrinsics") + missing = [ + name + for name in required + if getattr(prediction, name, None) is None + ] + if missing: + raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") + depth_frames = np.asarray(prediction.depth, dtype=np.float32) + confidence_frames = np.asarray(prediction.conf, dtype=np.float32) + if depth_frames.shape[0] != 1: + raise ValueError("Convert to Surface currently supports one image") + raw_depth = depth_frames[0] + confidence = confidence_frames[0] + intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32) + progress(2, "Unprojecting the depth grid") + points = unproject_camera_depth(raw_depth, intrinsics) + progress(3, "Preparing confidence and background data") + foreground_mask = np.ones(raw_depth.shape, dtype=bool) + if foreground_path is not None: + foreground_mask = foreground_mesh_mask( + foreground_path, + raw_depth.shape, + ) + confidence_percentiles = confidence_percentile_map( + raw_depth, + confidence, + ) + progress(4, "Building source triangles") + faces = mesh_triangle_faces( + points, + raw_depth, + confidence, + None, + False, + 0.0, + ) + if foreground_path is not None: + faces = remove_background_faces(faces, foreground_mask) + if not len(faces): + raise ValueError( + "DA3 output does not contain any valid surface faces" + ) + progress(5, "Preparing source UV and attributes") + depth_discontinuity = depth_discontinuity_ratio(raw_depth) + face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) + vertices, faces = compact_mesh( + points.reshape(-1, 3), + faces, + ) + height, width = raw_depth.shape + uv = project_camera_uv(vertices, intrinsics, (width, height)) + attributes = image_mesh_attributes( + vertices, + faces, + intrinsics, + confidence_percentiles, + face_depth_discontinuity, + ) + reference_depth = float(np.median(vertices[:, 2])) + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + progress(6, "Writing source surface data") + result_path = output_directory / "mesh.npz" + np.savez_compressed( + result_path, + vertices=vertices, + faces=faces, + uv=uv, + intrinsics=intrinsics, + image_size=np.asarray((width, height), dtype=np.int32), + reference_depth=np.asarray(reference_depth, dtype=np.float32), + **attributes, + ) + texture_path = output_directory / "texture.png" + texture_source = foreground_path or image_path + with Image.open(texture_source) as image: + image.convert("RGBA").save(texture_path, format="PNG", compress_level=1) + return result_path, texture_path + + +def depth_edge_mask(depth, threshold=12.0 / 255.0): + import numpy as np + + values = np.asarray(depth, dtype=np.float32) + edges = np.zeros(values.shape, dtype=bool) + for axis in (0, 1): + first = np.take(values, range(values.shape[axis] - 1), axis=axis) + second = np.take(values, range(1, values.shape[axis]), axis=axis) + valid = ( + np.isfinite(first) + & np.isfinite(second) + & (first > 0.0) + & (second > 0.0) + ) + scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6) + discontinuity = valid & (np.abs(first - second) / scale > threshold) + lower = [slice(None), slice(None)] + upper = [slice(None), slice(None)] + lower[axis] = slice(0, -1) + upper[axis] = slice(1, None) + edges[tuple(lower)] |= discontinuity + edges[tuple(upper)] |= discontinuity + return edges + + +def edge_culling_stretch_limit(culling): + factor = min(max(float(culling) / 100.0, 0.0001), 1.0) + return 1.0 / factor + + +def foreground_mesh_mask(foreground_path, depth_shape): + import numpy as np + from PIL import Image, ImageFilter + + height, width = depth_shape + with Image.open(foreground_path) as source: + alpha = source.convert("RGBA").getchannel("A") + alpha = alpha.resize( + (width, height), + resample=Image.Resampling.BILINEAR, + ) + alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75)) + return np.asarray(alpha) >= 128 + + +def confidence_percentile_threshold(depth, confidence, percentile): + import numpy as np + + if float(percentile) <= 0.0: + return None + valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence) + if not valid.any(): + raise ValueError("DA3 returned no valid depth confidence values") + return float(np.percentile(confidence[valid], percentile)) + + +def mesh_triangle_faces( + points, + depth, + confidence, + confidence_threshold, + filter_edges, + max_edge_length, + edge_threshold=12.0 / 255.0, + edge_stretch_limit=None, + vertex_mask=None, +): + import numpy as np + + height, width = depth.shape + indices = np.arange(height * width, dtype=np.int32).reshape(height, width) + quads = np.stack( + ( + indices[:-1, :-1], + indices[1:, :-1], + indices[1:, 1:], + indices[:-1, 1:], + ), + axis=-1, + ).reshape(-1, 4) + valid = ( + np.isfinite(points).all(axis=-1) + & np.isfinite(depth) + & (depth > 0.0) + ) + if confidence_threshold is not None: + valid &= np.isfinite(confidence) & ( + confidence >= confidence_threshold + ) + if filter_edges: + valid &= ~depth_edge_mask(depth, edge_threshold) + if vertex_mask is not None: + if np.shape(vertex_mask) != depth.shape: + raise ValueError("Foreground mask and depth dimensions do not match") + valid &= np.asarray(vertex_mask, dtype=bool) + triangles = np.concatenate( + (quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]), + axis=0, + ) + face_valid = valid.reshape(-1)[triangles].all(axis=1) + if edge_stretch_limit is not None: + face_valid &= triangle_edge_stretch_mask( + points, + triangles, + edge_stretch_limit, + ) + if max_edge_length > 0.0: + flat_points = points.reshape(-1, 3) + triangle_points = flat_points[triangles] + next_points = np.roll(triangle_points, -1, axis=1) + edge_lengths = np.linalg.norm( + triangle_points - next_points, + axis=-1, + ) + face_valid &= (edge_lengths <= max_edge_length).all(axis=1) + triangles = triangles[face_valid] + if not len(triangles): + return np.empty((0, 3), dtype=np.int32) + return triangles + + +def triangle_edge_stretch_mask(points, triangles, stretch_limit): + stretch = triangle_edge_stretch_values(points, triangles) + tolerance = float(stretch_limit) * (1.0 + 1e-5) + return (stretch <= tolerance).all(axis=1) + + +def export_textured_mesh_glb( + prediction, + image_paths, + output_directory, + confidence_percentile, + filter_edges, + max_edge_length, +): + import numpy as np + import trimesh + from PIL import Image + + required = ("depth", "conf", "intrinsics", "extrinsics") + missing = [name for name in required if getattr(prediction, name) is None] + if missing: + raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}") + depth = np.asarray(prediction.depth, dtype=np.float32) + confidence = np.asarray(prediction.conf, dtype=np.float32) + threshold = confidence_percentile_threshold( + depth, + confidence, + confidence_percentile, + ) + meshes = [] + referenced_points = [] + for index in range(depth.shape[0]): + points = unproject_depth_frame( + depth[index], + prediction.intrinsics[index], + prediction.extrinsics[index], + ) + faces = mesh_triangle_faces( + points, + depth[index], + confidence[index], + threshold, + filter_edges, + max_edge_length, + ) + if not len(faces): + continue + height, width = depth[index].shape + u = np.linspace(0.0, 1.0, width, dtype=np.float32) + v = np.linspace(1.0, 0.0, height, dtype=np.float32) + uu, vv = np.meshgrid(u, v) + uv = np.stack((uu, vv), axis=-1).reshape(-1, 2) + with Image.open(image_paths[index]) as source_image: + texture_image = source_image.convert("RGB").copy() + visual = trimesh.visual.texture.TextureVisuals( + uv=uv, + image=texture_image, + ) + mesh = trimesh.Trimesh( + vertices=points.reshape(-1, 3), + faces=faces, + visual=visual, + process=False, + ) + mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1) + mesh.vertex_attributes["o_confidence"] = confidence[ + index + ].reshape(-1) + mesh.vertex_attributes["o_confidence_percentile"] = ( + confidence_percentile_map( + depth[index], + confidence[index], + ).reshape(-1) + ) + mesh.remove_unreferenced_vertices() + meshes.append(mesh) + referenced_points.append(mesh.vertices) + if not meshes: + raise ValueError( + "Surface filters removed every face; lower Confidence Threshold or disable filters" + ) + + world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64) + if world_to_camera.shape == (3, 4): + homogeneous = np.eye(4, dtype=np.float64) + homogeneous[:3, :4] = world_to_camera + world_to_camera = homogeneous + cv_to_gltf = np.eye(4, dtype=np.float64) + cv_to_gltf[1, 1] = -1.0 + cv_to_gltf[2, 2] = -1.0 + alignment = cv_to_gltf @ world_to_camera + aligned_points = trimesh.transform_points( + np.concatenate(referenced_points, axis=0), + alignment, + ) + center = np.median(aligned_points, axis=0) + alignment[:3, 3] -= center + + scene = trimesh.Scene() + for index, mesh in enumerate(meshes): + mesh.apply_transform(alignment) + scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}") + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + output_path = output_directory / "mesh.glb" + scene.export(output_path) + return output_path diff --git a/src/snapform/server/media.py b/src/snapform/server/media.py new file mode 100644 index 0000000..aa751fa --- /dev/null +++ b/src/snapform/server/media.py @@ -0,0 +1,306 @@ +import shutil +import tempfile +import time +from pathlib import Path + + +IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} +VIDEO_SUFFIXES = { + ".mp4", + ".mov", + ".avi", + ".mkv", + ".webm", + ".m4v", + ".mts", + ".mpg", + ".mpeg", +} +SCENE_FILENAME = "scene.glb" + + +class TaskCancelled(RuntimeError): + pass + + +def check_cancelled(cancel_check): + if cancel_check is not None and cancel_check(): + raise TaskCancelled("Task cancelled") + + +def collect_images(directory): + directory = Path(directory) + return sorted( + path + for path in directory.iterdir() + if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES + ) + + +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, + av, + start_frame=1, + cancel_check=None, +): + output_directory = Path(output_directory) + output_directory.mkdir(parents=True, exist_ok=True) + try: + container = av.open(str(video_path)) + except Exception as error: + raise RuntimeError(f"Unable to open video: {video_path}") from error + image_paths = [] + skipped = 0 + try: + stream = container.streams.video[0] + for frame in container.decode(stream): + check_cancelled(cancel_check) + if skipped < start_frame - 1: + skipped += 1 + continue + image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg" + frame.to_image().save(image_path, format="JPEG") + image_paths.append(image_path) + if len(image_paths) >= maximum_frames: + break + finally: + container.close() + + 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, + 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 + if suffix not in VIDEO_SUFFIXES: + supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES)) + raise RuntimeError(f"Unsupported input type; expected one of: {supported}") + + import av + + temporary_directory = Path( + tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root) + ) + try: + images = extract_video_frames( + input_path, + temporary_directory, + video_frames, + av, + start_frame=start_frame, + cancel_check=cancel_check, + ) + except Exception: + shutil.rmtree(temporary_directory, ignore_errors=True) + raise + 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" + destination = ( + Path(output_directory) / f"frame.{index + 1:04d}{suffix}" + ) + if Path(image_path).resolve() == destination.resolve(): + continue + shutil.copy2(image_path, destination) + + +def extract_frames(context, parameters): + """Copy selected source frames into the current job directory.""" + context.directory.mkdir(parents=True, exist_ok=True) + images, temporary_directory = prepare_input( + parameters["input"], + parameters.get("video_frames", 16), + context.directory, + start_frame=parameters.get("start_frame", 1), + cancel_check=context.check_cancelled, + ) + try: + copy_frames(images, context.directory) + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) + message = f"Extracted {len(images)} frame(s)" + context.progress(1.0, message) + return {"frames": len(images)} + + +def exported_paths(output_directory, export_format): + output_directory = Path(output_directory) + if export_format == "glb": + scene_path = output_directory / SCENE_FILENAME + return [scene_path] if scene_path.is_file() else [] + if export_format == "gs_ply": + return sorted((output_directory / "gs_ply").glob("*.ply")) + if export_format == "depth_vis": + return sorted((output_directory / "depth_vis").glob("*.jpg")) + if export_format == "mesh": + result_path = output_directory / "mesh.glb" + return [result_path] if result_path.is_file() else [] + if export_format == "mesh_data": + result_path = output_directory / "mesh.npz" + return [result_path] if result_path.is_file() else [] + if export_format == "mini_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 [] + + +def npz_is_complete(path): + """Return True only when the npz archive can be fully read.""" + try: + import numpy as np + + with np.load(path, allow_pickle=False) as archive: + return "depth" in archive.files + except Exception: + return False + + +def wait_for_exported_paths( + output_directory, + export_format, + timeout=30.0, + cancel_check=None, +): + # DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整， + # 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取 + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + check_cancelled(cancel_check) + result_paths = exported_paths(output_directory, export_format) + if result_paths: + if ( + export_format in ("mini_npz", "npz") + and not npz_is_complete(result_paths[0]) + ): + time.sleep(0.1) + continue + return result_paths + time.sleep(0.05) + return [] + + +def move_exported_file( + source, + destination, + timeout=10.0, + cancel_check=None, +): + """Move an async export after Windows releases its file handle.""" + source = Path(source) + destination = Path(destination) + deadline = time.monotonic() + timeout + while True: + check_cancelled(cancel_check) + try: + source.replace(destination) + return destination + except OSError as error: + is_file_lock = isinstance(error, PermissionError) or getattr( + error, + "winerror", + None, + ) in {32, 33} + if not is_file_lock or time.monotonic() >= deadline: + raise + time.sleep(0.1) + + +def flatten_exported_paths( + output_directory, + result_paths, + cancel_check=None, +): + """Move files produced in exporter subdirectories to the job root.""" + output_directory = Path(output_directory) + flattened_paths = [] + for source in result_paths: + source = Path(source) + check_cancelled(cancel_check) + if source.parent == output_directory: + flattened_paths.append(source) + continue + destination = output_directory / source.name + if destination.exists(): + raise RuntimeError(f"Export filename already exists: {destination.name}") + flattened_paths.append( + move_exported_file( + source, + destination, + cancel_check=cancel_check, + ) + ) + for directory in sorted( + {path.parent for path in map(Path, result_paths)}, + key=lambda path: len(path.parts), + reverse=True, + ): + while directory != output_directory: + try: + directory.rmdir() + except OSError: + break + directory = directory.parent + return flattened_paths diff --git a/src/snapform/models.py b/src/snapform/server/model_catalog.py similarity index 62% rename from src/snapform/models.py rename to src/snapform/server/model_catalog.py index cdabc10..fffc66a 100644 --- a/src/snapform/models.py +++ b/src/snapform/server/model_catalog.py @@ -153,125 +153,6 @@ MODEL_NUMBERS = { } -_ENUM_ITEMS_CACHE = {} -_DEPTH_PLANE_ENUM_ITEMS_CACHE = {} -_SURFACE_ENUM_ITEMS_CACHE = {} - - -def enum_items(supports_gs=None, show_download_status=False): - model_ready = None - if show_download_status: - from .client import model_ready - - models = [ - model - for model in MODELS.values() - if model.enabled - and (supports_gs is None or model.supports_gs == supports_gs) - ] - download_status = ( - tuple(model_ready(model.key) for model in models) - if model_ready is not None - else () - ) - cache_key = ( - supports_gs, - show_download_status, - download_status, - ) - cached_items = _ENUM_ITEMS_CACHE.get(cache_key) - if cached_items is not None: - return cached_items - - items = [] - for index, model in enumerate(models): - label = model.label - if download_status and not download_status[index]: - label = f"{label} (Missing)" - items.append( - ( - model.key, - label, - f"{model.description}; license: {model.license_name}", - MODEL_NUMBERS[model.key], - ) - ) - - _ENUM_ITEMS_CACHE[cache_key] = items - return items - - -def model_enum_items(_owner, _context): - return enum_items(show_download_status=True) - - -def gs_model_enum_items(_owner, _context): - return enum_items(supports_gs=True, show_download_status=True) - - -def moge2_model_enum_items(_owner, _context): - from .client import model_ready - - return [ - ( - model.key, - model.label if model_ready(model.key) else f"{model.label} (Missing)", - f"{model.description}; license: {model.license_name}", - index, - ) - for index, model in enumerate(MOGE2_MODELS.values()) - ] - - -def depth_plane_model_enum_items(_owner, _context): - from .client import model_ready - - models = [ - *MOGE2_MODELS.values(), - *(model for model in MODELS.values() if model.enabled), - ] - ready = tuple(model_ready(model.key) for model in models) - cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready) - if cached is not None: - return cached - items = [] - for index, model in enumerate(models): - label = model.label if ready[index] else f"{model.label} (Missing)" - items.append( - ( - model.key, - label, - f"{model.description}; license: {model.license_name}", - index, - ) - ) - items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items))) - _DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items - return items - - -def surface_model_enum_items(owner, context): - depth_items = depth_plane_model_enum_items(owner, context) - cache_key = id(depth_items) - cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key) - if cached is not None: - return cached - items = [ - item - for item in depth_items - if item[0] != "ORIGINAL" - ] - _SURFACE_ENUM_ITEMS_CACHE[cache_key] = items - return items - - -def selected_model_key(settings, export_format=None): - selected_format = export_format or settings.export_format - if selected_format == "gs_ply": - return settings.gs_model - return settings.model - - def get_model(model_key): try: return MODELS[model_key] @@ -289,3 +170,20 @@ def get_downloadable_model(model_key): def model_directory_name(model_key): """Return the model's folder name inside the shared models directory.""" return get_downloadable_model(model_key).directory_name + + +def model_record(model, ready=False): + """Return public model metadata exposed by the Models Resource.""" + return { + "key": model.key, + "label": model.label, + "description": model.description, + "license": model.license_name, + "supports_gs": model.supports_gs, + "enabled": model.enabled, + "family": "moge2" if model.key.startswith("MOGE2_") else ( + "ben2" if model.key == BEN2_MODEL.key else "da3" + ), + "ready": bool(ready), + } + diff --git a/src/snapform/server/model_download.py b/src/snapform/server/model_download.py new file mode 100644 index 0000000..9bc48a3 --- /dev/null +++ b/src/snapform/server/model_download.py @@ -0,0 +1,444 @@ +import hashlib +import http.client +import os +import threading +import time +from pathlib import Path +from urllib.parse import quote +from urllib.request import ( + HTTPSHandler, + ProxyHandler, + Request, + build_opener, + urlopen, +) + +HF_ENDPOINT = "https://huggingface.co" +R2_ENDPOINT = "https://models.omoolab.xyz" +CONNECT_TIMEOUT = 3.0 +DOWNLOAD_CHUNK_SIZE = 256 * 1024 + + +class _ShortConnectHTTPSConnection(http.client.HTTPSConnection): + """Use a short timeout for connection setup, without limiting downloads.""" + + def connect(self): + try: + self.timeout = CONNECT_TIMEOUT + super().connect() + finally: + self.timeout = None + self.sock.settimeout(None) + + +class _ShortConnectHTTPSHandler(HTTPSHandler): + def https_open(self, request): + return self.do_open(_ShortConnectHTTPSConnection, request) + + +_download_opener = build_opener(ProxyHandler(), _ShortConnectHTTPSHandler()) + + +class DownloadProgressReporter: + def __init__( + self, + progress_callback, + repository, + clock=time.monotonic, + start=0.05, + span=0.90, + label=None, + ): + self.progress_callback = progress_callback + self.repository = repository + self.clock = clock + self.lock = threading.Lock() + self.start = start + self.span = span + self.label = label + self.progress = start + self.last_write_time = 0.0 + self.started_at = self.clock() + + def report(self, current, total, rate=None): + if not total or current <= 0: + return + + factor = min(max(current / total, 0.0), 1.0) + progress = self.start + factor * self.span + now = self.clock() + with self.lock: + if progress <= self.progress: + return + if now - self.last_write_time < 0.20 and factor < 1.0: + return + self.progress = progress + self.last_write_time = now + effective_rate = rate + elapsed = now - self.started_at + if not effective_rate and elapsed > 0: + effective_rate = current / elapsed + speed = format_rate(effective_rate) if effective_rate else "" + prefix = self.label or self.repository + message = f"Downloading {prefix}" + if speed: + message = f"{message} {speed}" + self.progress_callback( + progress, + message, + ) + + +def create_download_progress_class( + progress_callback, + repository, + tqdm_type=None, + clock=time.monotonic, + start=0.05, + span=0.90, + label=None, +): + if tqdm_type is None: + from tqdm.auto import tqdm as tqdm_type + + reporter = DownloadProgressReporter( + progress_callback, + repository, + clock, + start, + span, + label, + ) + + class DownloadProgress(tqdm_type): + def __init__(self, *args, **kwargs): + description = str(kwargs.get("desc", "")) + # huggingface_hub 1.x 的字节进度条描述是 "Downloading (incomplete total...)"， + # 用前缀匹配而非精确匹配，否则真实下载进度会被忽略。 + self.snapform_reports_progress = ( + description.startswith("Downloading") + or description.startswith("Fetching ") + ) + super().__init__(*args, **kwargs) + + def update(self, amount=1): + result = super().update(amount) + if self.snapform_reports_progress: + reporter.report(self.n, self.total, self._download_rate()) + return result + + def refresh(self, *args, **kwargs): + result = super().refresh(*args, **kwargs) + if self.snapform_reports_progress: + reporter.report(self.n, self.total, self._download_rate()) + return result + + def _download_rate(self): + try: + return getattr(self, "format_dict", {}).get("rate") + except Exception: + return None + + return DownloadProgress + + +def progress_range(args): + start = getattr(args, "progress_start", None) + end = getattr(args, "progress_end", None) + if start is not None and end is not None: + return float(start), float(end) - float(start) + return 0.05, 0.90 + + +def download_model(args): + label = getattr(args, "label", None) or args.repository + model_dir = Path(args.model_dir) + model_dir.mkdir(parents=True, exist_ok=True) + start, span = progress_range(args) + r2_directory = getattr(args, "r2_directory", "") + r2_files = getattr(args, "r2_file", []) + if r2_directory and r2_files: + _download_declared_or_fail( + args, + model_dir, + label, + start, + span, + ) + return + _download_snapshot_or_fail(args, model_dir, label, start, span) + + +def _download_declared_or_fail(args, model_dir, label, start, span): + """Download declared files from Hugging Face, then the R2 mirror.""" + sources = ( + (HF_ENDPOINT, download_from_hf, f"Downloading {label}"), + ( + R2_ENDPOINT, + download_from_r2, + f"Downloading {label} (mirror)", + ), + ) + failures = [] + for source, downloader, message in sources: + args.progress(start, message) + try: + downloader( + args, + model_dir, + start=start, + span=span, + label=label, + ) + _verify_model_complete( + model_dir, + getattr(args, "required_pattern", None), + ) + except Exception as error: + cancel_check = getattr(args, "cancel_check", None) + if cancel_check is not None: + cancel_check() + failures.append((source, error)) + continue + args.progress(start + span, f"{label} ready") + return + detail = "; ".join( + f"{source}: {error}" for source, error in failures + ) + raise RuntimeError(f"All download sources failed: {detail}") + + +def _download_snapshot_or_fail(args, model_dir, label, start, span): + """Fallback download for repositories without declared files, via huggingface_hub.""" + failures = [] + try: + from huggingface_hub import snapshot_download + except Exception as error: + snapshot_download = None + failures.append((HF_ENDPOINT, error)) + progress_class = None + if snapshot_download: + progress_class = create_download_progress_class( + args.progress, + args.repository, + start=start, + span=span, + label=label, + ) + for endpoint in download_endpoints() if snapshot_download else (): + os.environ["HF_ENDPOINT"] = endpoint + # snapshot_download 先拉取文件清单并校验 ETag，还没有字节下载， + # 用 Checking 提示当前阶段，避免进度条看起来卡住。 + message = f"Checking {label} files" + if endpoint != HF_ENDPOINT: + message = f"{message} (mirror)" + args.progress(start, message) + 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, + getattr(args, "required_pattern", None), + ) + except Exception as error: + cancel_check = getattr(args, "cancel_check", None) + if cancel_check is not None: + cancel_check() + failures.append((endpoint, error)) + continue + args.progress(start + span, f"{label} ready") + return + detail = "; ".join( + f"{source}: {error}" for source, error in failures + ) + raise RuntimeError(f"All download sources failed: {detail}") + + +def download_endpoints(): + """Return the configured Hugging Face endpoint or official source.""" + endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT + return [endpoint.rstrip("/")] + + +def _verify_model_complete(model_dir, required_patterns=None): + """Refuse to report success when the required model files are missing.""" + model_dir = Path(model_dir) + required_patterns = required_patterns or ( + "config.json", + "*.safetensors", + ) + missing = [ + pattern + for pattern in required_patterns + if not any(model_dir.glob(pattern)) + ] + if missing: + raise RuntimeError( + "Model download reported success but files are missing: " + + ", ".join(missing) + ) + + +def _download_file_from_url( + url, + destination, + expected_size, + expected_checksum, + progress=None, + retries=1, + cancel_check=None, +): + """Download and verify one model file atomically from a direct URL.""" + request = Request(url, headers={"User-Agent": "SnapForm"}) + opener = _download_opener.open if request.type == "https" else urlopen + last_error = None + for _attempt in range(retries): + temporary = destination.with_name(f"{destination.name}.part") + try: + checksum = hashlib.sha256() + downloaded = 0 + with opener(request) as response: + with open(temporary, "wb") as output: + while True: + if cancel_check is not None: + cancel_check() + chunk = response.read(DOWNLOAD_CHUNK_SIZE) + if not chunk: + break + output.write(chunk) + checksum.update(chunk) + downloaded += len(chunk) + if progress is not None: + progress(len(chunk)) + if downloaded != expected_size: + raise RuntimeError( + f"Size mismatch for {destination.name}: " + f"{downloaded} != {expected_size}" + ) + if checksum.hexdigest().lower() != expected_checksum.lower(): + raise RuntimeError(f"SHA-256 mismatch for {destination.name}") + temporary.replace(destination) + return + except Exception as error: + if cancel_check is not None: + cancel_check() + last_error = error + try: + temporary.unlink() + except OSError: + pass + raise last_error + + +def r2_file_url(directory, filename): + return f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}" + + +def hf_file_url(repository, filename): + """Hugging Face resolve URL for a direct, metadata-free download.""" + endpoint = (os.environ.get("HF_ENDPOINT") or HF_ENDPOINT).rstrip("/") + return f"{endpoint}/{repository}/resolve/main/{quote(filename)}" + + +def _parse_declared_files(args): + selected = [] + for value in args.r2_file: + filename, size, checksum = value.split("|", 2) + selected.append((filename, int(size), checksum)) + return selected + + +def _download_declared_files( + args, + model_dir, + url_builder, + start, + span, + label, + retries=1, +): + """Download each declared file through url_builder with aggregated progress.""" + repository = args.repository + selected = _parse_declared_files(args) + reporter = DownloadProgressReporter( + args.progress, + repository, + start=start, + span=span, + label=label or repository, + ) + total_size = sum(size for _path, size, _checksum in selected) + downloaded = 0 + + def on_chunk(count): + nonlocal downloaded + downloaded += count + reporter.report(downloaded, total_size) + + for path, size, checksum in selected: + destination = Path(model_dir) / path + destination.parent.mkdir(parents=True, exist_ok=True) + _download_file_from_url( + url_builder(path), + destination, + size, + checksum, + progress=on_chunk, + retries=retries, + cancel_check=getattr(args, "cancel_check", None), + ) + + +def download_from_r2( + args, + model_dir, + start=0.05, + span=0.90, + label=None, +): + """Download the declared model files from OmooLab R2.""" + _download_declared_files( + args, + model_dir, + lambda filename: r2_file_url(args.r2_directory, filename), + start, + span, + f"{label} (mirror)" if label else "mirror", + ) + + +def download_from_hf( + args, + model_dir, + start=0.05, + span=0.90, + label=None, +): + """Download the declared model files directly from Hugging Face resolve URLs.""" + _download_declared_files( + args, + model_dir, + lambda filename: hf_file_url(args.repository, filename), + start, + span, + label, + ) + + +def format_rate(rate): + size = float(rate) + if size <= 0: + return "" + units = ("B/s", "KB/s", "MB/s", "GB/s") + unit_index = 0 + while size >= 1024 and unit_index < len(units) - 1: + size /= 1024 + unit_index += 1 + if unit_index == 0: + return f"{size:.0f} {units[unit_index]}" + return f"{size:.1f} {units[unit_index]}" diff --git a/src/snapform/ai_service/model_manager.py b/src/snapform/server/model_manager.py similarity index 73% rename from src/snapform/ai_service/model_manager.py rename to src/snapform/server/model_manager.py index c21f17a..586ed53 100644 --- a/src/snapform/ai_service/model_manager.py +++ b/src/snapform/server/model_manager.py @@ -3,12 +3,12 @@ import time from pathlib import Path from types import SimpleNamespace -import jobs -import models as model_catalog +from . import model_catalog +from .model_download import download_model class ModelSessionCache: - """Cache ONNX sessions shared by jobs in one service process.""" + """Cache ONNX sessions shared by Jobs in one Server process.""" def __init__(self): self.ben2_model = None @@ -27,7 +27,7 @@ class ModelSessionCache: return self._load_ben2(model_dir, requested_device, key) def _load_ben2(self, model_dir, requested_device, key): - from ben2_onnx import create_session + from .onnx_ben2 import create_session started = time.perf_counter() model = create_session(model_dir, requested_device) @@ -47,15 +47,15 @@ class ModelSessionCache: return self._load_da3(model_dir, requested_device, key) def _load_da3(self, model_dir, requested_device, key): - from da3_onnx import create_session + from .onnx_da3 import create_session model = create_session(model_dir, requested_device) self.da3_model = model self.da3_key = key return model - def get_moge_onnx(self, model_dir, requested_device): - from moge_onnx import create_session + def get_moge2(self, model_dir, requested_device): + from .onnx_moge2 import create_session key = (str(Path(model_dir).resolve()), requested_device, "onnx") if self.da3_model is not None: @@ -112,7 +112,27 @@ class ModelManager: def directory(self, model_key): model = model_catalog.get_downloadable_model(model_key) - return self.models_directory / model.directory_name + directory = self.models_directory / model.directory_name + old_moge_directory = self.models_directory / "moge2_base_normal" + if ( + model_key == "MOGE2_VITB_NORMAL" + and not directory.is_dir() + and old_moge_directory.is_dir() + ): + old_moge_directory.rename(directory) + legacy_directory = self.models_directory / model_key.lower() + if ( + not directory.is_dir() + and legacy_directory.is_dir() + and legacy_directory != directory + ): + legacy_directory.rename(directory) + return directory + + def ready(self, model_key): + model = model_catalog.get_downloadable_model(model_key) + directory = self.directory(model_key) + return all(any(directory.glob(pattern)) for pattern in model.ready_patterns) def get_ben2(self, model_dir, device): return self.sessions.get_ben2(model_dir, device) @@ -120,11 +140,15 @@ class ModelManager: def get_da3(self, model_dir, device): return self.sessions.get_da3(model_dir, device) - def get_moge_onnx(self, model_dir, device): - return self.sessions.get_moge_onnx(model_dir, device) + def get_moge2(self, model_dir, device): + return self.sessions.get_moge2(model_dir, device) def snapshot(self): - return {"loaded": self.sessions.snapshot()} + catalog = [ + model_catalog.model_record(model, self.ready(model.key)) + for model in model_catalog.DOWNLOADABLE_MODELS.values() + ] + return {"loaded": self.sessions.snapshot(), "catalog": catalog} def clear(self): self.sessions.close() @@ -137,7 +161,7 @@ class ModelManager: args = SimpleNamespace( repository=model.repository, model_dir=str(self.directory(model_key)), - status=context.update, + progress=context.progress, cancel_check=context.check_cancelled, label=model.label, allow_pattern=list(model.download_patterns), @@ -149,5 +173,5 @@ class ModelManager: ], ) context.check_cancelled() - jobs.download_model(args) + download_model(args) context.check_cancelled() diff --git a/src/snapform/server/moge2.py b/src/snapform/server/moge2.py new file mode 100644 index 0000000..723281b --- /dev/null +++ b/src/snapform/server/moge2.py @@ -0,0 +1,304 @@ +import shutil +from pathlib import Path + +from .background import infer_foreground, infer_foreground_frames +from .geometry import ( + depth_discontinuity_ratio, + depth_edge_mask, + foreground_mesh_mask, + mesh_triangle_faces, +) +from .media import copy_frames, prepare_input + + +def _save_moge_preview(output_directory, index, depth, normal, mask): + import numpy as np + from PIL import Image + + suffix = "" if index is None else f".{index + 1:04d}" + valid = mask & np.isfinite(depth) & (depth > 0) + preview = np.zeros(depth.shape, dtype=np.uint8) + if np.any(valid): + low, high = np.percentile(depth[valid], (2.0, 98.0)) + span = max(float(high - low), 1e-6) + preview[valid] = np.clip( + (depth[valid] - low) / span * 255.0, + 0.0, + 255.0, + ).astype(np.uint8) + Image.fromarray(preview, mode="L").save( + output_directory / f"depth{suffix}.png" + ) + normal_rgb = np.clip( + (normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0) + * 127.5, + 0.0, + 255.0, + ).astype(np.uint8) + Image.fromarray(normal_rgb, mode="RGB").save( + output_directory / f"normal{suffix}.png" + ) + Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save( + output_directory / f"mask{suffix}.png" + ) + + +def export_moge_debug_mesh( + output_directory, + points, + depth, + mask, + normal, + image, + export_format, +): + import numpy as np + import trimesh + + clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04) + faces = mesh_triangle_faces( + points, + depth, + np.ones(depth.shape, dtype=np.float32), + None, + False, + 0.0, + vertex_mask=clean_mask, + ) + vertices = points.reshape(-1, 3) + colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8) + vertex_normals = normal.reshape(-1, 3) + vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32) + vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32) + mesh = trimesh.Trimesh( + vertices=vertices, + faces=faces, + vertex_colors=colors, + vertex_normals=vertex_normals, + process=False, + ) + mesh.remove_unreferenced_vertices() + filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply" + mesh.export(Path(output_directory) / filename) + + +def export_moge_surface_archive( + output_path, + points, + depth, + mask, + normal, + intrinsics, + foreground_mask=None, +): + import numpy as np + + height, width = depth.shape + finite_mask = np.isfinite(depth) & (depth > 0.0) + finite_mask &= np.isfinite(points).all(axis=-1) + if foreground_mask is not None: + finite_mask &= foreground_mask + depth_discontinuity = depth_discontinuity_ratio(depth) + faces = mesh_triangle_faces( + points, + depth, + np.ones(depth.shape, dtype=np.float32), + None, + False, + 0.0, + vertex_mask=finite_mask, + ) + vertices = points.reshape(-1, 3) + u = np.linspace(0.0, 1.0, width, dtype=np.float32) + v = np.linspace(1.0, 0.0, height, dtype=np.float32) + uu, vv = np.meshgrid(u, v) + vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2) + pixel_intrinsics = intrinsics.copy() + pixel_intrinsics[0] *= width + pixel_intrinsics[1] *= height + pixel_intrinsics[2] = (0.0, 0.0, 1.0) + point_depth = depth.reshape(-1).astype(np.float32) + face_mask = mask.reshape(-1)[faces].all(axis=1) + face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1) + masked_depth = point_depth[mask.reshape(-1).astype(bool)] + reference_depth = np.median(masked_depth if len(masked_depth) else point_depth) + np.savez( + output_path, + vertices=vertices.astype(np.float32), + faces=faces.astype(np.int32), + uv=vertex_uvs.astype(np.float32), + intrinsics=pixel_intrinsics.astype(np.float32), + image_size=np.array((width, height), dtype=np.int32), + reference_depth=np.float32(reference_depth), + point_depth=point_depth, + normal_map=normal.astype(np.float32), + face_mask=face_mask, + face_depth_discontinuity=face_depth_discontinuity, + ) + + +def generate_moge2(context, parameters): + """Run MoGe-2 inference and return the artifact used by the requesting Job.""" + import numpy as np + from PIL import Image + from .onnx_moge2 import ( + infer as infer_moge2, + provider_label, + ) + + context.progress(0.02, "Starting MoGe-2") + output_directory = context.directory + output_directory.mkdir(parents=True, exist_ok=True) + model_manager = context.resource("model_manager") + model_directory = model_manager.directory(parameters["model"]) + device = parameters["device"] + images, temporary_directory = prepare_input( + parameters["input"], + parameters.get("video_frames", 16), + output_directory, + start_frame=parameters.get("start_frame", 1), + cancel_check=context.check_cancelled, + ) + try: + model = model_manager.get_moge2(model_directory, device) + context.progress( + 0.08, + f"Loading MoGe-2 on {provider_label(model)}", + ) + + depths = [] + normals = [] + masks = [] + points = [] + intrinsics = [] + source_images = [] + for index, image_path in enumerate(images): + context.check_cancelled() + context.progress( + 0.15 + 0.68 * index / len(images), + f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})", + ) + with Image.open(image_path) as opened: + image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0 + source_images.append(image) + prediction = infer_moge2( + model, + image, + int(parameters["resolution_level"]), + ) + depth = prediction["depth"] + normal = prediction["normal"] + mask = prediction["mask"].astype(bool) + point = prediction["points"] + intrinsic = prediction["intrinsics"] + depths.append(depth) + normals.append(normal) + masks.append(mask) + intrinsics.append(intrinsic) + if bool( + parameters.get("include_points", False) + or parameters.get("surface_mesh", False) + ): + points.append(point) + debug_format = parameters.get("export_format", "") + if not debug_format or debug_format == "maps": + _save_moge_preview( + output_directory, + None if len(images) == 1 else index, + depth, + normal, + mask, + ) + + shapes = {array.shape for array in depths} + if len(shapes) != 1: + raise RuntimeError("MoGe-2 frame sizes must match") + result = { + "depth": np.stack(depths).astype(np.float32), + "normal": np.stack(normals).astype(np.float32), + "mask": np.stack(masks), + "intrinsics": np.stack(intrinsics).astype(np.float32), + } + if points: + result["points"] = np.stack(points).astype(np.float32) + debug_format = parameters.get("export_format", "") + if debug_format in {"glb", "ply"} and len(images) == 1: + context.progress( + 0.84, + f"Exporting MoGe-2 {debug_format.upper()}", + ) + export_moge_debug_mesh( + output_directory, + points[0], + depths[0], + masks[0], + normals[0], + source_images[0], + debug_format, + ) + + foreground_mask = None + remove_background_model = parameters.get("remove_background_model", "") + if remove_background_model: + remove_background_model_dir = model_manager.directory( + remove_background_model + ) + context.progress( 0.87, "Removing image background") + if len(images) == 1: + foreground, _, _ = infer_foreground( + images[0], + remove_background_model_dir, + device, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + foreground_path = output_directory / "foreground.png" + foreground.save(foreground_path) + if points: + foreground_mask = foreground_mesh_mask( + foreground_path, + result["depth"][0].shape, + ) + else: + infer_foreground_frames( + images, + remove_background_model_dir, + device, + output_directory, + context.progress, + 0.87, + 0.08, + model_cache=model_manager, + cancel_check=context.check_cancelled, + ) + if bool(parameters.get("surface_mesh", False)): + if len(images) != 1 or not points: + raise RuntimeError("MoGe-2 Surface requires one point map") + context.progress( 0.94, "Building official MoGe-2 surface") + export_moge_surface_archive( + output_directory / "mesh.npz", + points[0], + result["depth"][0], + result["mask"][0], + result["normal"][0], + result["intrinsics"][0], + foreground_mask=( + foreground_mask + if remove_background_model + else None + ), + ) + np.savez(output_directory / "results.npz", **result) + copy_frames(images, output_directory) + context.progress( + 1.0, + f"Exported MoGe-2 data for {len(images)} frame(s)", + ) + if context.job_type == "generate-moge2-surface": + return {"surface": "mesh.npz"} + if context.job_type == "generate-moge2-depth": + return {"depth": "results.npz"} + return {"directory": "."} + finally: + if temporary_directory is not None: + shutil.rmtree(temporary_directory, ignore_errors=True) diff --git a/src/snapform/ai_service/ben2_onnx.py b/src/snapform/server/onnx_ben2.py similarity index 99% rename from src/snapform/ai_service/ben2_onnx.py rename to src/snapform/server/onnx_ben2.py index 9ab9803..3132e57 100644 --- a/src/snapform/ai_service/ben2_onnx.py +++ b/src/snapform/server/onnx_ben2.py @@ -48,3 +48,4 @@ def infer(session, image): foreground = image.convert("RGBA") foreground.putalpha(mask) return foreground + diff --git a/src/snapform/ai_service/da3_onnx.py b/src/snapform/server/onnx_da3.py similarity index 99% rename from src/snapform/ai_service/da3_onnx.py rename to src/snapform/server/onnx_da3.py index cc89a61..ca9ef7a 100644 --- a/src/snapform/ai_service/da3_onnx.py +++ b/src/snapform/server/onnx_da3.py @@ -240,3 +240,4 @@ def infer( processed_images=processed_images, aux={}, ) + diff --git a/src/snapform/ai_service/moge_onnx.py b/src/snapform/server/onnx_moge2.py similarity index 99% rename from src/snapform/ai_service/moge_onnx.py rename to src/snapform/server/onnx_moge2.py index 5e5d390..b11a323 100644 --- a/src/snapform/ai_service/moge_onnx.py +++ b/src/snapform/server/onnx_moge2.py @@ -163,3 +163,4 @@ def infer(session, image, resolution_level): raw_output = dict(zip(names, values)) raw_output["metric_scale"] = raw_output.pop(scale_name) return postprocess(raw_output) + diff --git a/src/snapform/ai_service/onnx_runtime.py b/src/snapform/server/onnx_runtime.py similarity index 100% rename from src/snapform/ai_service/onnx_runtime.py rename to src/snapform/server/onnx_runtime.py diff --git a/src/snapform/status_bar.py b/src/snapform/status_bar.py deleted file mode 100644 index 5b3206a..0000000 --- a/src/snapform/status_bar.py +++ /dev/null @@ -1,55 +0,0 @@ -from .operators import active_job - - -def draw_status_bar(self, context): - scene = getattr(context, "scene", None) - job = active_job() - if job is None or scene is None: - return - - settings = getattr(scene, "snapform_settings", None) - if settings is None: - return - - progress = min(max(settings.progress, 0.0), 1.0) - row = self.layout.row(align=True) - progress_row = row.row(align=True) - progress_row.ui_units_x = 18.0 - progress_row.progress( - factor=progress, - type="BAR", - text=settings.status, - ) - row.operator( - "snapform.cancel_job", - text="", - icon="X", - ) - - -def tag_ui_regions(context): - window_manager = getattr(context, "window_manager", None) - if window_manager is None: - return - - for window in window_manager.windows: - screen = getattr(window, "screen", None) - if screen is None: - continue - for area in screen.areas: - if area.type in {"STATUSBAR", "VIEW_3D"}: - area.tag_redraw() - for region in getattr(area, "regions", ()): - region.tag_redraw() - - -def redraw_ui(context, force=False): - tag_ui_regions(context) - if not force: - return - workspace = getattr(context, "workspace", None) - if workspace is not None: - try: - workspace.status_text_set_internal(None) - except (AttributeError, RuntimeError, TypeError): - pass diff --git a/tests/test_addon.py b/tests/test_addon.py index 9a17289..546f365 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -1,15 +1,13 @@ import importlib +import json import sys +import tempfile import threading -import time import unittest from pathlib import Path -from tempfile import TemporaryDirectory from types import ModuleType, SimpleNamespace from unittest.mock import patch -import numpy as np - PROJECT_ROOT = Path(__file__).resolve().parents[1] SOURCE_ROOT = PROJECT_ROOT / "src" @@ -19,6 +17,11 @@ class AddonTest(unittest.TestCase): @classmethod def setUpClass(cls): cls.remove_modules() + + def append_status_bar(draw): + draw._owner = "snapform" + cls.status_bar_draws.append(draw) + fake_bpy = ModuleType("bpy") fake_bpy.props = SimpleNamespace( BoolProperty=lambda **options: options, @@ -36,7 +39,7 @@ class AddonTest(unittest.TestCase): PropertyGroup=object, Scene=type("Scene", (), {}), STATUSBAR_HT_header=SimpleNamespace( - append=lambda draw: cls.status_bar_draws.append(draw), + append=append_status_bar, remove=lambda draw: cls.status_bar_draws.remove(draw), ), VIEW3D_MT_object_context_menu=SimpleNamespace( @@ -81,2442 +84,440 @@ class AddonTest(unittest.TestCase): if module_name == "snapform" or module_name.startswith("snapform."): sys.modules.pop(module_name) + def setUp(self): + preferences = SimpleNamespace( + storage_root="", + device="CUDA", + ) + self.fake_bpy.context.preferences.addons = { + "snapform": SimpleNamespace(preferences=preferences) + } + def tearDown(self): + self.addon.runtime.close_active() self.registered.clear() self.status_bar_draws.clear() self.context_menu_draws.clear() self.outliner_menu_draws.clear() - self.addon.operators.close_active_job() 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 - def test_uses_extension_manifest_instead_of_legacy_bl_info(self): - self.assertFalse(hasattr(self.addon, "bl_info")) - - def test_registers_and_unregisters_extension_types(self): - self.addon.register() - - self.assertEqual(self.registered, list(self.addon.CLASSES)) - self.assertTrue( - hasattr(self.fake_bpy.types.Scene, "snapform_settings") - ) - self.assertEqual( - self.status_bar_draws, - [self.addon.draw_status_bar], - ) - self.assertEqual( - self.context_menu_draws, - [self.addon.draw_image_context_menu], - ) - self.assertEqual( - self.outliner_menu_draws, - [self.addon.draw_image_context_menu], - ) - - self.addon.unregister() - - self.assertEqual(self.registered, []) - self.assertFalse( - hasattr(self.fake_bpy.types.Scene, "snapform_settings") - ) - self.assertEqual(self.status_bar_draws, []) - self.assertEqual(self.context_menu_draws, []) - self.assertEqual(self.outliner_menu_draws, []) - - def test_registers_depth_plane_operator_without_cache_rebuild(self): - self.assertIn( - self.addon.ConvertImageToDepthPlane, - self.addon.CLASSES, - ) - self.assertIn(self.addon.BuildDepthPlane, self.addon.CLASSES) - self.assertEqual( - self.addon.BuildDepthPlane.bl_options, - {"INTERNAL"}, - ) - self.assertIn( - "snapform_color_frames_dir", - self.addon.BuildDepthPlane.__annotations__, - ) - self.assertIn( - "snapform_depth_frames_dir", - self.addon.BuildDepthPlane.__annotations__, - ) - self.assertIn("execute", self.addon.BuildDepthPlane.__dict__) - self.assertFalse(hasattr(self.addon, "ImportDepthPlane")) - - def test_registers_perspective_mesh_operators(self): - self.assertIn(self.addon.ConvertImageToMesh, self.addon.CLASSES) - self.assertIn(self.addon.BuildImageMesh, self.addon.CLASSES) - self.assertEqual(self.addon.BuildImageMesh.bl_options, {"INTERNAL"}) - self.assertIn(self.addon.SnapFormImageMenu, self.addon.CLASSES) - - def test_registered_job_operators_own_their_blender_properties(self): - expected = { - self.addon.operators.GenerateDa3Surface: { - "snapform_input_path", - "snapform_use_ray_pose", - "snapform_process_res", - "snapform_mesh_create_camera", - }, - self.addon.operators.GenerateMoge2Surface: { - "snapform_input_path", - "snapform_moge_resolution_level", - "snapform_mesh_create_camera", - }, - self.addon.operators.GenerateDa3Depth: { - "snapform_depth_data_name", - "snapform_use_ray_pose", - "snapform_process_res", - }, - self.addon.operators.GenerateMoge2Depth: { - "snapform_depth_data_name", - "snapform_moge_resolution_level", - }, - } - - for operator, properties in expected.items(): - with self.subTest(operator=operator.__name__): - self.assertTrue(properties <= set(operator.__annotations__)) - - def test_image_context_menu_groups_actions_and_separates_background(self): - calls = [] - - class Layout: - operator_context = None - - def operator(self, identifier, **_options): - calls.append(identifier) - - def separator(self): - calls.append("separator") - - menu = self.addon.SnapFormImageMenu() - menu.layout = Layout() - with patch.object(self.addon.client, "env_installed", return_value=True): - menu.draw(None) - - self.assertEqual( - calls, - [ - "snapform.convert_image_to_depth_plane", - "snapform.convert_image_to_mesh", - "separator", - "snapform.remove_image_background", - ], - ) - - def test_image_context_menu_only_installs_env_when_missing(self): - calls = [] - - class Operator: - job_type = None - - class Layout: - operator_context = None - - def operator(self, identifier, **options): - operator = Operator() - calls.append((identifier, options, operator)) - return operator - - def separator(self): - calls.append("separator") - - menu = self.addon.SnapFormImageMenu() - menu.layout = Layout() - with patch.object(self.addon.client, "env_installed", return_value=False): - menu.draw(None) - - self.assertEqual(len(calls), 1) - identifier, options, operator = calls[0] - self.assertEqual(identifier, "snapform.install_environment") - self.assertEqual(options["text"], "Install Environment") - - def test_image_mesh_reference_plane_matches_image_empty_bounds(self): - source_object = SimpleNamespace( - data=SimpleNamespace(size=(401, 201)), - empty_display_size=4.0, - empty_image_offset=(-0.25, -0.75), - ) - intrinsics = np.asarray( - [[200.0, 0.0, 200.0], [0.0, 200.0, 100.0], [0.0, 0.0, 1.0]], - dtype=np.float32, - ) - camera_vertices = np.asarray( - [ - [-2.0, -1.0, 2.0], - [2.0, -1.0, 2.0], - [-2.0, 1.0, 2.0], - [2.0, 1.0, 2.0], - ], - dtype=np.float32, - ) - - vertices, camera = self.addon.image_mesh.empty_anchored_vertices( - source_object, - camera_vertices, - intrinsics, - (401, 201), - 2.0, - ) - - x_min, x_max, y_min, y_max = self.addon.image_depth.plane_bounds( - source_object - ) - np.testing.assert_allclose( - vertices, - [ - [x_min, y_max, 0.0], - [x_max, y_max, 0.0], - [x_min, y_min, 0.0], - [x_max, y_min, 0.0], - ], - ) - np.testing.assert_allclose( - camera["location"], - ((x_min + x_max) * 0.5, (y_min + y_max) * 0.5, 2.0024922), - ) - - def test_mesh_conversion_dialog_only_keeps_model_and_camera_options(self): - annotations = self.addon.ConvertImageToMesh.__annotations__ - - for property_name in ( - "snapform_confidence_percentile", - "snapform_edge_culling", - "snapform_face_reduction", - "snapform_depth_smoothing", - "snapform_postprocess_method", + def test_snapform_has_no_runtime_boundary_modules(self): + for name in ( + "client.py", + "environment.py", + "storage.py", + "models.py", + "status_bar.py", ): - with self.subTest(property_name=property_name): - self.assertNotIn(property_name, annotations) - self.assertFalse(annotations["snapform_create_camera"]["default"]) - self.assertFalse(annotations["snapform_unshaded"]["default"]) - self.assertEqual(annotations["snapform_unshaded"]["name"], "Unshaded") - self.assertFalse(annotations["snapform_remove_background"]["default"]) - self.assertFalse(annotations["snapform_keep_original"]["default"]) - self.assertEqual( - annotations["snapform_process_res_factor"]["default"], - 100.0, - ) - self.assertEqual(annotations["snapform_process_res_factor"]["min"], 1.0) - self.assertEqual(annotations["snapform_process_res_factor"]["max"], 100.0) - self.assertEqual( - annotations["snapform_process_res_factor"]["subtype"], - "PERCENTAGE", - ) - self.assertEqual( - annotations["snapform_model"]["default"], - 0, - ) + self.assertFalse((SOURCE_ROOT / "snapform" / name).exists(), name) + self.assertFalse(hasattr(self.addon, "client")) - def test_image_mesh_uses_o_prefixed_attributes_and_bundled_nodes(self): - image_mesh = self.addon.image_mesh - names = set(image_mesh.POINT_ATTRIBUTES) | set( - image_mesh.FACE_ATTRIBUTES - ) + def test_job_request_is_not_part_of_blendjob(self): + blendjob = importlib.import_module("snapform.blendjob") + self.assertFalse(hasattr(blendjob, "JobRequest")) - self.assertTrue(all(name.startswith("o_") for name in names)) - self.assertTrue(image_mesh.node_asset_path().is_file()) - self.assertEqual(image_mesh.NODE_GROUP_NAME, "O DA3 Surface") - self.assertEqual(image_mesh.NODE_ASSET_NAME, "O_SnapForm.blend") - - def test_depth_plane_uses_bundled_nodes_and_material_attributes(self): - image_depth = self.addon.image_depth - - self.assertTrue(image_depth.node_asset_path().is_file()) - self.assertEqual( - image_depth.DEPTH_PLANE_NODE_GROUP_NAME, - "O DA3 Depth Plane", - ) - self.assertEqual( - image_depth.MATERIAL_DISPLACEMENT_ATTRIBUTE, - "o_material_displacement", - ) - self.assertEqual( - image_depth.BACKGROUND_DEPTH_ATTRIBUTE, - "o_background_depth", - ) - - def test_image_mesh_material_depth_scale_matches_fitted_geometry(self): - mesh_object = SimpleNamespace(scale=SimpleNamespace(z=-2.0)) - - self.assertEqual( - self.addon.image_mesh.material_depth_scale( - mesh_object, - {"scale_z": 0.25}, - ), - 0.5, - ) - - def test_image_mesh_default_resolution_uses_double_reduction_estimate(self): - vertices = np.asarray( - ((0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (0.0, 1.0, 0.0)), - dtype=np.float32, - ) - faces = np.asarray(((0, 1, 2),), dtype=np.int32) - edge_lengths = np.asarray((1.0, 1.0, np.sqrt(2.0))) - - self.assertAlmostEqual( - self.addon.image_mesh.default_remesh_resolution(vertices, faces), - float(np.median(edge_lengths) * np.sqrt(10.0) * 2.0), - ) - - def test_registers_image_background_operators(self): - self.assertIn(self.addon.RemoveImageBackground, self.addon.CLASSES) - self.assertIn(self.addon.ApplyRemovedBackground, self.addon.CLASSES) - self.assertIn(self.addon.BuildOriginalDepthPlane, self.addon.CLASSES) - self.assertEqual( - self.addon.ApplyRemovedBackground.bl_options, - {"INTERNAL"}, - ) - - def test_background_result_creates_new_packed_image(self): - original_image = SimpleNamespace( - name="img.png", - use_fake_user=False, - ) - transform = object() - source_object = SimpleNamespace( - data=original_image, - matrix_world=transform, - empty_display_size=2.5, - empty_image_offset=(-0.25, -0.75), - ) - pack_calls = [] - result_image = SimpleNamespace( - name="foreground.png", - pack=lambda: pack_calls.append(True), - ) - load_calls = [] - remove_calls = [] - self.fake_bpy.data = SimpleNamespace( - images=SimpleNamespace( - load=lambda *args, **options: load_calls.append( - (args, options) - ) - or result_image, - remove=remove_calls.append, - ), - ) - - with TemporaryDirectory() as directory: - output_path = Path(directory) / "foreground.png" - output_path.touch() - created = self.addon.image_depth.create_removed_background_image( - source_object, - output_path, - ) - - self.assertIs(created, result_image) - self.assertIs(source_object.data, result_image) - self.assertIsNot(source_object.data, original_image) - self.assertTrue(original_image.use_fake_user) - self.assertIs(source_object.matrix_world, transform) - self.assertEqual(source_object.empty_display_size, 2.5) - self.assertEqual( - source_object.empty_image_offset, - (-0.25, -0.75), - ) - self.assertEqual(result_image.name, "img.foreground") - self.assertEqual(result_image.alpha_mode, "CHANNEL_PACKED") - self.assertEqual(pack_calls, [True]) - self.assertEqual(remove_calls, []) - self.assertEqual( - load_calls[0][1], - {"check_existing": False}, - ) - - def test_background_result_name_handles_numbered_format_suffix(self): - source_object = SimpleNamespace( - data=SimpleNamespace(name="img.png.001"), - ) - - self.assertEqual( - self.addon.image_depth.removed_background_data_name( - source_object - ), - "img.001.foreground", - ) - - def test_depth_background_image_does_not_replace_source_empty(self): - original_image = SimpleNamespace( - name="img.png", - use_fake_user=False, - ) - source_object = SimpleNamespace(data=original_image) - pack_calls = [] - result_image = SimpleNamespace( - name="foreground.png", - pack=lambda: pack_calls.append(True), - ) - self.fake_bpy.data = SimpleNamespace( - images=SimpleNamespace( - load=lambda *_args, **_options: result_image, - remove=lambda _image: None, - ), - ) - - with TemporaryDirectory() as directory: - output_path = Path(directory) / "foreground.png" - output_path.touch() - created = self.addon.image_depth.create_removed_background_image( - source_object, - output_path, - assign_to_source=False, - ) - - 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.foreground") - self.assertEqual(result_image.alpha_mode, "CHANNEL_PACKED") - 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(result_image.alpha_mode, "CHANNEL_PACKED") - 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=SimpleNamespace( - snapform_settings=SimpleNamespace(video_frames=16) - ), - ) - calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - process_image_background=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, True), - ): - result = operator.execute(context) - - self.assertEqual(result, {"FINISHED"}) - self.assertNotIn("job_type", calls[0]) - self.assertEqual(calls[0]["snapform_input_path"], str(input_path)) - self.assertEqual( - calls[0]["snapform_source_object_name"], - "Reference", - ) - self.assertTrue(calls[0]["snapform_delete_input"]) - - def test_original_depth_background_starts_independent_job(self): - operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_model = "ORIGINAL" - operator.snapform_unshaded = False - operator.snapform_remove_background = True - operator.snapform_keep_original = True - source_object = SimpleNamespace( - name="Reference", - data=SimpleNamespace(name="img.png"), - ) - context = SimpleNamespace(object=source_object, scene=object()) - calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - generate_original_background=lambda **options: calls.append(options) - or {"RUNNING_MODAL"}, - ), - ) - - 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.assertNotIn("job_type", calls[0]) - self.assertNotIn("snapform_background_depth_distance", calls[0]) - - def test_original_depth_animation_extracts_frames_and_background(self): - operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_model = "ORIGINAL" - operator.snapform_unshaded = False - operator.snapform_remove_background = True - operator.snapform_keep_original = True - 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( - generate_original_background=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.assertNotIn("job_type", calls[0]) - 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_original_depth_without_background_uses_frame_extraction(self): - operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_model = "ORIGINAL" - operator.snapform_unshaded = False - operator.snapform_remove_background = False - operator.snapform_keep_original = False - source_object = SimpleNamespace( - name="Reference", - data=SimpleNamespace(name="img.png"), - ) - context = SimpleNamespace(object=source_object, scene=object()) - calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - extract_original_frames=lambda **options: calls.append(options) - or {"RUNNING_MODAL"}, - ), - ) - - 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.assertNotIn("job_type", calls[0]) - self.assertTrue(calls[0]["snapform_delete_source"]) - - def test_conversion_dialog_exposes_requested_settings(self): - operator = self.addon.ConvertImageToDepthPlane - - self.assertEqual( - operator.__annotations__["snapform_model"]["default"], - 0, - ) - self.assertFalse(operator.__annotations__["snapform_use_ray_pose"]["default"]) - self.assertEqual( - operator.__annotations__["snapform_process_res_factor"]["default"], - 100.0, - ) - self.assertEqual( - operator.__annotations__["snapform_process_res_factor"]["min"], - 1.0, - ) - self.assertEqual( - operator.__annotations__["snapform_process_res_factor"]["max"], - 100.0, - ) - self.assertFalse( - operator.__annotations__["snapform_remove_background"]["default"] - ) - self.assertFalse( - operator.__annotations__["snapform_keep_original"]["default"] - ) - self.assertFalse( - operator.__annotations__["snapform_unshaded"]["default"] - ) - self.assertEqual( - set(operator.__annotations__), - { - "snapform_unshaded", - "snapform_model", - "snapform_use_ray_pose", - "snapform_process_res_factor", - "snapform_moge_resolution_level", - "snapform_remove_background", - "snapform_keep_original", - }, - ) - self.assertEqual( - operator.__annotations__["snapform_model"]["name"], - "Depth From", - ) - self.assertEqual( - operator.__annotations__["snapform_unshaded"]["name"], - "Unshaded", - ) - model_items_callback = operator.__annotations__["snapform_model"]["items"] - with patch.object( - self.addon.client, - "model_ready", - return_value=True, + def test_job_operators_are_flat_and_use_plain_properties(self): + operators = self.addon.operators + for name in ( + "SnapFormJobMixin", + "ImageInputJobOperator", + "GeometryImageJobOperator", + "Da3ImageJobOperator", + "Moge2ImageJobOperator", + "BackgroundRequestMixin", ): - model_items = model_items_callback(None, None) + self.assertFalse(hasattr(operators, name), name) + for operator in operators.JOB_OPERATOR_CLASSES: + for name in getattr(operator, "__annotations__", {}): + self.assertFalse(name.startswith("snapform_"), (operator, name)) + + def test_concrete_job_operator_is_the_blender_operator(self): + operator = self.addon.operators.GenerateMoge2Depth + self.assertEqual(operator.bl_idname, "snapform.generate_moge2_depth") + self.assertIn(self.addon.job_runtime.JobOperatorBase, operator.__mro__) + self.assertIn(self.fake_bpy.types.Operator, operator.__mro__) + + def test_default_request_sends_every_operator_property(self): + class Example(self.addon.job_runtime.JobOperatorBase): + value: object() + enabled: object() + + operator = Example() + operator.value = 42 + operator.enabled = True self.assertEqual( - [item[0] for item in model_items], - [ - "MOGE2_VITS_NORMAL", - "MOGE2_VITB_NORMAL", - "MOGE2_VITL_NORMAL", - "SMALL", - "BASE", - "ORIGINAL", - ], + operator.request(object()), + {"value": 42, "enabled": True}, ) - def test_conversion_dialog_leaves_last_properties_to_blender(self): - operator = self.addon.ConvertImageToDepthPlane() - 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, - ), + def test_overridden_request_purely_replaces_default_properties(self): + class Example(self.addon.job_runtime.JobOperatorBase): + value: object() + + def request(self, context): + return {"replacement": context.value} + + operator = Example() + operator.value = 42 + self.assertEqual( + operator.request(SimpleNamespace(value=7)), + {"replacement": 7}, ) - result = operator.invoke(context, None) - - self.assertIs(result, operator) - - 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) + def test_runtime_submits_job_type_and_plain_dict(self): calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - process_image_background=lambda **options: calls.append(options) - or {"RUNNING_MODAL"}, - ), - ) - operator = self.addon.RemoveImageBackground() + submitted = threading.Event() - 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) + class Controller: + def submit(self, job_type, parameters): + calls.append((job_type, parameters)) + submitted.set() + return {"job_id": "job", "directory": "."} - self.assertEqual(result, {"FINISHED"}) - self.assertEqual(scene.snapform_settings.video_frames, 24) + def cancel(self, _job_id): + pass - def test_conversion_execute_uses_background_mask_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_model = "BASE" - operator.snapform_use_ray_pose = True - operator.snapform_process_res_factor = 50.0 - operator.snapform_unshaded = False - operator.snapform_remove_background = True - operator.snapform_keep_original = True - calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - generate_da3_depth=lambda **options: calls.append(options) - or {"RUNNING_MODAL"}, - ), - ) + def mark_job_complete(self, _job_id): + pass - 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.assertNotIn("job_type", calls[0]) - self.assertTrue(calls[0]["snapform_remove_background"]) - self.assertNotIn("snapform_background_depth_distance", calls[0]) - - def test_conversion_dialog_marks_missing_models(self): - items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[ - "snapform_model" - ]["items"] - with patch.object( - self.addon.client, - "model_ready", - side_effect=lambda key: key in {"SMALL", "BASE"}, - ): - items = items_callback(None, None) - - self.assertEqual( - [item[0] for item in items], - [ - "MOGE2_VITS_NORMAL", - "MOGE2_VITB_NORMAL", - "MOGE2_VITL_NORMAL", - "SMALL", - "BASE", - "ORIGINAL", - ], - ) - labels = {item[0]: item[1] for item in items} - self.assertEqual(labels["BASE"], "DA3 Base") - self.assertEqual(labels["MOGE2_VITB_NORMAL"], "MoGe-2 ViT-B Normal (Missing)") - - def test_update_status_formats_stage_progress(self): - operator = self.addon.operators.ModalJobOperatorBase() - - self.assertEqual( - operator.format_job_status( - { - "progress": 0.5, - "message": "2 downloaded", - "stage": 3, - "stages": 8, - "stage_label": "pytorch", - }, - "fallback", - ), - "3/8 2 downloaded", - ) - self.assertEqual( - operator.format_job_status( - { - "progress": 0.0, - "message": "", - "stage": 6, - "stages": 8, - "stage_label": "da3 small", - }, - "fallback", - ), - "6/8 da3 small", - ) - self.assertEqual( - operator.format_job_status( - { - "progress": 0.5, - "message": "Estimating depth", - }, - "fallback", - ), - "Estimating depth", - ) - self.assertEqual( - operator.format_job_status( - { - "progress": 0.25, - "message": "Downloading packages", - "stage": 2, - "stages": 3, - "stage_label": "PyTorch", - "eta_seconds": 125, - }, - "fallback", - ), - "2/3 Downloading packages", - ) - self.assertEqual( - operator.format_job_status( - { - "progress": 0.5, - "message": "", - "stage": 2, - "stages": 3, - "stage_label": "PyTorch", - }, - "2/3 PyTorch | Downloading packages", - ), - "2/3 PyTorch", - ) - - def test_modal_job_operator_only_requires_type_and_parameters(self): - controller = object() - - class ExampleJob(self.addon.operators.ModalJobOperatorBase): - controller_factory = staticmethod(lambda: controller) + class Example(self.addon.job_runtime.JobOperatorBase): job_type = "example" + value: object() - def job_parameters(self, context): - return {"value": context.value} - - operator = ExampleJob() - request = operator.create_job_request(SimpleNamespace(value=42)) - - self.assertIs(operator.get_service_controller(), controller) - self.assertEqual(request.job_type, "example") - self.assertEqual(request.parameters, {"value": 42}) - - def test_update_status_resets_progress_when_stage_changes(self): - operator = self.addon.operators.SnapFormJobOperator() - settings = SimpleNamespace(progress=0.9, status="previous") - context = SimpleNamespace( - scene=SimpleNamespace(snapform_settings=settings) - ) - job = SimpleNamespace(progress_stage=1, progress=0.9) - - with patch.object(operator, "redraw_job_ui"): - operator.update_from_job_status( - context, - job, - { - "progress": 0.1, - "message": "Downloading model", - "stage": 2, - "stages": 4, - "stage_label": "MoGe-2", - }, - ) - - self.assertEqual(job.progress_stage, 2) - self.assertEqual(settings.progress, 0.1) - - def test_model_options_mark_only_models_that_are_not_downloaded(self): - items_callback = self.addon.SnapFormSettings.__annotations__["model"][ - "items" - ] - with patch.object( - self.addon.client, - "model_ready", - side_effect=lambda key: key == "BASE", - ): - items = items_callback(None, None) - - labels = {item[0]: item[1] for item in items} - self.assertEqual(labels["BASE"], "DA3 Base") - self.assertEqual(labels["SMALL"], "DA3 Small (Missing)") - - def test_packed_image_is_exported_without_unpacking(self): - operator = self.addon.ConvertImageToDepthPlane() - operator.snapform_model = "GIANT_1_1" - operator.snapform_use_ray_pose = True - operator.snapform_process_res_factor = 50.0 - operator.snapform_unshaded = True - operator.snapform_remove_background = True - operator.snapform_keep_original = False - image = SimpleNamespace( - name="img.png", - filepath="", - packed_file=object(), - size=(1920, 1080), - ) - context = SimpleNamespace( - object=SimpleNamespace(name="Reference", data=image), - scene=object(), - ) - calls = [] - self.fake_bpy.ops = SimpleNamespace( - snapform=SimpleNamespace( - generate_da3_depth=lambda **options: calls.append(options) - or {"RUNNING_MODAL"}, - ), - ) - - with TemporaryDirectory() as directory: - packed_path = Path(directory) / "packed.png" - packed_path.touch() - with patch.object( - self.addon.image_depth, - "save_packed_image", - return_value=packed_path, - ) as save_packed_image, patch.object( - self.addon.client, - "model_ready", - return_value=True, - ): - result = operator.execute(context) - - self.assertEqual(result, {"FINISHED"}) - save_packed_image.assert_called_once_with(image, context.scene) - self.assertEqual(calls[0]["snapform_input_path"], str(packed_path)) - self.assertEqual(calls[0]["snapform_process_res"], 966) - self.assertEqual(calls[0]["snapform_model_key"], "GIANT_1_1") - self.assertTrue(calls[0]["snapform_remove_background"]) - self.assertTrue(calls[0]["snapform_unshaded"]) - self.assertEqual(calls[0]["snapform_depth_data_name"], "img.depth") - self.assertTrue(calls[0]["snapform_delete_input"]) - - def test_packed_image_export_writes_a_png_copy(self): - class PackedImage: - def save_render(self, filepath, scene): - self.filepath = filepath - self.scene = scene - self.file_format = scene.render.image_settings.file_format - Path(filepath).write_bytes(b"packed image") - - image = PackedImage() - image_settings = SimpleNamespace(file_format="OPEN_EXR") - scene = SimpleNamespace( - render=SimpleNamespace(image_settings=image_settings), - ) - with TemporaryDirectory() as directory: - output_path = Path(directory) / "input.png" - with patch.object( - self.addon.client, - "temporary_input_path", - return_value=output_path, - ): - result = self.addon.image_depth.save_packed_image(image, scene) - - self.assertEqual(result, output_path) - self.assertEqual(output_path.read_bytes(), b"packed image") - self.assertEqual(image.filepath, str(output_path)) - self.assertIs(image.scene, scene) - self.assertEqual(image.file_format, "PNG") - self.assertEqual(image_settings.file_format, "OPEN_EXR") - - def test_temporary_packed_input_is_removed_when_job_closes(self): - with TemporaryDirectory() as directory: - input_path = Path(directory) / "packed.png" - input_path.touch() - data = self.addon.operators.SnapFormJobData( - kind="test", - temporary_input_path=input_path, - ) - - self.addon.operators.SnapFormJobOperator().cleanup_job_data(data) - - self.assertFalse(input_path.exists()) - - def test_temporary_frame_directory_is_removed_when_job_closes(self): - with TemporaryDirectory() as directory: - root = Path(directory) - temporary_directory = root / "inputs" / "frames" - temporary_directory.mkdir(parents=True) - (temporary_directory / "frame-0000.jpg").touch() - data = self.addon.operators.SnapFormJobData( - kind="test", - temporary_input_path=temporary_directory, - ) - - self.addon.operators.SnapFormJobOperator().cleanup_job_data(data) - - self.assertFalse(temporary_directory.exists()) - - 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.client, - "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", - export_format="gs_ply", - model="BASE", - gs_model="GIANT_1_1", - ) - operator = self.addon.operators.GenerateDa3Depth() - operator.snapform_source_object_name = "Reference" - operator.snapform_process_res = 336 - - 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", - } - for snapform_model, snapform_ray_pose in ( - ("SMALL", True), - ("BASE", False), - ("GIANT_1_1", False), - ("NESTED_GIANT_LARGE_1_1", True), - ): - with self.subTest(model=snapform_model), patch.object( - self.addon.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ) as model_ready, patch.object( - self.addon.client, - "generate_da3_depth", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as reconstruct: - operator.snapform_model_key = snapform_model - operator.snapform_use_ray_pose = snapform_ray_pose - command = operator.create_request(settings) - - self.assertIsInstance( - command, - self.addon.operators.JobRequest, - ) - model_ready.assert_called_once_with(snapform_model) - self.assertEqual(reconstruct.call_args.args[1], input_path) - self.assertEqual(reconstruct.call_args.args[2], snapform_model) - self.assertEqual( - reconstruct.call_args.kwargs["process_res"], - 336, - ) - self.assertEqual( - reconstruct.call_args.kwargs["use_ray_pose"], - snapform_ray_pose, - ) - - def test_image_depth_background_mask_requires_ben2(self): - settings = SimpleNamespace( - input_path="debug-input.png", - export_format="mini_npz", - model="BASE", - gs_model="GIANT_1_1", - ) - operator = self.addon.operators.GenerateDa3Depth() - operator.snapform_model_key = "BASE" - operator.snapform_use_ray_pose = True - operator.snapform_process_res = 1024 - operator.snapform_remove_background = True - - with TemporaryDirectory() as directory: - input_path = Path(directory) / "input.png" - input_path.touch() - operator.snapform_input_path = str(input_path) - with patch.object( - self.addon.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - side_effect=lambda key: key == "BASE", - ): - with self.assertRaisesRegex(RuntimeError, "BEN2 Base"): - operator.create_request(settings) - - def test_image_depth_background_option_reaches_service_request(self): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace( - input_path="input.png", - 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", - ) - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - side_effect=lambda key: Path(key.lower()), - ): - request = importlib.import_module("snapform.client").generate_da3_depth( - settings, - Path("input.png"), - "BASE", - process_res=504, - use_ray_pose=False, - remove_background=True, - ) - - self.assertEqual( - request.payload["parameters"]["remove_background_model"], - "BEN2_BASE", - ) - self.assertNotIn( - "background_depth_distance", - request.payload["parameters"], - ) - - def test_service_busy_snapshot_is_released_for_completed_local_job(self): - client = self.addon.client - controller = client.service_controller() - previous_snapshot = controller.snapshot - controller.snapshot = { - "state": "BUSY", - "busy": True, - "active_job": {"job_id": "depth-job"}, - } - try: - controller.mark_job_complete("depth-job") - - self.assertEqual(client.service_status()["state"], "READY") - self.assertFalse(client.service_busy()) - finally: - controller.snapshot = previous_snapshot - - def test_job_submission_starts_without_blocking_execute(self): - operator = self.addon.operators.GenerateDa3Depth() - operator.snapform_source_object_name = "Reference" - operator.snapform_depth_data_name = "img.depth" - operator.snapform_process_res = 1024 - operator.snapform_delete_source = False - operator.snapform_delete_input = False - operator.snapform_unshaded = False - started = threading.Event() - release = threading.Event() - - def submit(_request): - started.set() - release.wait(timeout=2.0) + def controller(self, _runtime): + return Controller() + timers = [] window_manager = SimpleNamespace( - event_timer_add=lambda *_args, **_kwargs: object(), - modal_handler_add=lambda *_args: None, + progress_begin=lambda *_args: None, + progress_update=lambda _value: None, + event_timer_add=lambda *_args, **_kwargs: timers.append(object()) or timers[-1], + event_timer_remove=lambda _timer: None, + modal_handler_add=lambda _operator: None, ) context = SimpleNamespace( - scene=SimpleNamespace( - snapform_settings=SimpleNamespace(status="", progress=0.0) - ), window_manager=window_manager, window=object(), ) - request = self.addon.operators.JobRequest({"job_id": "job"}) - try: - with patch.object( - operator, - "create_job_request", - return_value=request, - ), patch.object( - self.addon.client.service_controller(), - "submit", - side_effect=submit, - ), patch.object( - self.addon.client, - "service_log_path", - return_value=Path("host.log"), - ), patch.object( - self.addon.client, - "storage_root", - return_value=Path("."), - ), patch.object( - self.addon.client, - "creation_flags", - return_value=0, - ), patch.object(operator, "redraw_job_ui"): - before = time.perf_counter() - result = operator.execute(context) - elapsed = time.perf_counter() - before + operator = Example() + operator.value = 9 - self.assertEqual(result, {"RUNNING_MODAL"}) - self.assertLess(elapsed, 0.2) - self.assertTrue(started.wait(timeout=1.0)) - finally: - release.set() - job = self.addon.operators.active_job() - if job is not None and job.start_thread is not None: - job.start_thread.join(timeout=2.0) - self.addon.operators.close_active_job() + self.assertEqual(operator.execute(context), {"RUNNING_MODAL"}) + self.assertTrue(submitted.wait(1.0)) + self.assertEqual(calls, [("example", {"value": 9})]) - def test_preferences_use_package_identifier(self): - self.assertEqual( - self.addon.SnapFormPreferences.bl_idname, - "snapform", - ) - - def test_default_storage_is_shared_in_user_home(self): - client = importlib.import_module("snapform.client") - - self.assertEqual(client.DEFAULT_STORAGE_ROOT.name, ".snapform") - self.assertEqual( - self.addon.SnapFormPreferences.__annotations__["storage_root"][ - "default" - ], - str(client.DEFAULT_STORAGE_ROOT), - ) - - def test_ai_service_uses_loopback_address(self): - client = importlib.import_module("snapform.client") - - self.assertEqual(client.SERVICE_ADDRESS, "127.0.0.1") - - def test_closing_blender_terminates_service_it_started(self): - client = importlib.import_module("snapform.client") - controller = client.service_controller() - with patch.object(controller, "detach") as detach: - client.detach_service_client() - detach.assert_called_once_with() - - def test_process_res_method_defaults_to_upper_bound_resize(self): - process_res = self.addon.SnapFormSettings.__annotations__[ - "process_res" - ] - process_res_method = ( - self.addon.SnapFormSettings.__annotations__[ - "process_res_method" - ] - ) - - self.assertEqual(process_res["default"], 504) - self.assertEqual(process_res_method["default"], "upper_bound_resize") - - def test_reconstruct_request_preserves_da3_parameter_names(self): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace( - input_path="input.png", - export_format="gs_ply", - model="GIANT_1_1", - gs_model="NESTED_GIANT_LARGE_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="CUDA", - use_ray_pose=False, - ref_view_strategy="saddle_balanced", - ) - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - return_value=Path("model"), - ): - request = importlib.import_module("snapform.client").debug_da3( - settings, - Path("input.png"), - "NESTED_GIANT_LARGE_1_1", - ) - - parameters = request.payload["parameters"] - self.assertEqual(request.payload["job_type"], "debug-da3") - self.assertEqual(parameters["inference_mode"], "joint") - self.assertEqual(parameters["export_format"], "gs_ply") - for parameter in ( - "process_res", - "process_res_method", - "conf_thresh_percentile", - "num_max_points", - "show_cameras", - "gs_views_interval", - "use_ray_pose", - "ref_view_strategy", - ): - with self.subTest(parameter=parameter): - self.assertIn(parameter, parameters) - - def test_remove_background_request_passes_video_frames(self): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace(device="AUTO", video_frames=24) - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - return_value=Path("ben2_base"), - ): - request = importlib.import_module("snapform.client").remove_background( - settings, - Path("clip.mp4"), - "BEN2_BASE", - start_frame=70, - ) - - parameters = request.payload["parameters"] - self.assertEqual(parameters["video_frames"], 24) - self.assertEqual(parameters["start_frame"], 70) - self.assertNotIn("background_depth_distance", parameters) - - def test_extract_frames_request_passes_frame_range(self): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace(device="AUTO", video_frames=24) - - request = importlib.import_module("snapform.client").extract_frames( - settings, - Path("clip.mp4"), - start_frame=70, - ) - - parameters = request.payload["parameters"] - self.assertEqual(request.payload["job_type"], "extract-frames") - self.assertEqual(parameters["input"], str(Path("clip.mp4"))) - self.assertNotIn("output_dir", parameters) - self.assertEqual(parameters["video_frames"], 24) - self.assertEqual(parameters["start_frame"], 70) - - 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.operators.GenerateDa3Depth() - 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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "generate_da3_depth", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as reconstruct: - operator.create_request(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): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace( - input_path="input.png", - 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", - ) - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - return_value=Path("large-model"), - ) as snapform_model_directory: - request = importlib.import_module("snapform.client").generate_da3_depth( - settings, - Path("input.png"), - "BASE", - process_res=252, - use_ray_pose=False, - ) - - snapform_model_directory.assert_not_called() - parameters = request.payload["parameters"] - self.assertEqual(parameters["model"], "BASE") - self.assertEqual(parameters["process_res"], 252) - self.assertEqual( - parameters["process_res_method"], - "upper_bound_resize", - ) - self.assertEqual(parameters["use_ray_pose"], 0) - - def test_reconstruct_request_enables_ben2_for_image_mesh(self): - client = importlib.import_module("snapform.client") - settings = SimpleNamespace( - input_path="input.png", - 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", - ) - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - side_effect=lambda key: Path(key.lower()), - ): - request = importlib.import_module("snapform.client").generate_da3_surface( - settings, - Path("input.png"), - "BASE", - process_res=504, - use_ray_pose=False, - remove_background=True, - ) - - parameters = request.payload["parameters"] - self.assertNotIn("mesh_confidence_percentile", parameters) - self.assertNotIn("mesh_edge_culling", parameters) - self.assertNotIn("mesh_face_reduction", parameters) - self.assertNotIn("mesh_depth_smoothing", parameters) - self.assertNotIn("mesh_postprocess_method", parameters) - self.assertEqual( - parameters["remove_background_model"], - "BEN2_BASE", - ) - - def test_export_format_property_includes_all_da3_formats(self): - export_format = self.addon.SnapFormSettings.__annotations__[ - "export_format" - ] - identifiers = [item[0] for item in export_format["items"]] - - self.assertEqual( - identifiers, - ["glb", "gs_ply", "depth_vis", "mini_npz", "mesh"], - ) - self.assertEqual(export_format["default"], "glb") - self.assertEqual(export_format["name"], "Output Type") - self.assertEqual( - [item[1] for item in export_format["items"]], - [ - "Point Cloud (GLB)", - "Gaussian Splats (PLY)", - "Depth Preview (Images)", - "Depth Data (NPZ)", - "Textured Surface (GLB)", - ], - ) - - def test_mesh_is_a_da3_output_type(self): - settings = SimpleNamespace( - input_path="", - export_format="mesh", - 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="CPU", - use_ray_pose=False, - ref_view_strategy="saddle_balanced", - ) - operator = self.addon.operators.DebugDa3() - - with TemporaryDirectory() as directory: - input_path = Path(directory) / "input.png" - input_path.touch() - settings.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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "debug_da3", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as reconstruct: - operator.create_request(settings) - - self.assertEqual(reconstruct.call_args.args[1], input_path) - self.assertEqual(reconstruct.call_args.args[2], "BASE") - - def test_confidence_threshold_is_a_bounded_percentage(self): - confidence = self.addon.SnapFormSettings.__annotations__[ - "conf_thresh_percentile" - ] - - self.assertEqual(confidence["name"], "Confidence Threshold") - self.assertEqual(confidence["min"], 0.0) - self.assertEqual(confidence["max"], 100.0) - self.assertEqual(confidence["subtype"], "PERCENTAGE") - - def test_gs_model_property_only_includes_gaussian_models(self): - gs_model = self.addon.SnapFormSettings.__annotations__[ - "gs_model" - ] - identifiers = [item[0] for item in gs_model["items"](None, None)] - - self.assertEqual( - identifiers, - [], - ) - - def test_device_enums_list_only_supported_devices(self): - client = importlib.import_module("snapform.client") - with patch.object( - client, - "supported_devices", - return_value=["cuda", "cpu"], - ): - items = client.device_enum_items(None, None) - - self.assertEqual([item[0] for item in items], ["CUDA", "CPU"]) - self.assertEqual(items[0][2], "NVIDIA CUDA") - - def test_device_enums_default_to_the_fastest_supported_device(self): - for class_type in ( - self.addon.SnapFormSettings, - self.addon.SnapFormPreferences, - ): - with self.subTest(class_type=class_type.__name__): - self.assertEqual( - class_type.__annotations__["device"]["default"], - 0, - ) - - def test_preferences_expose_production_device(self): - preferences = self.addon.SnapFormPreferences - - self.assertIn("device", preferences.__annotations__) - items_callback = preferences.__annotations__["device"]["items"] - with patch.object( - self.addon.client, - "supported_devices", - return_value=["cpu"], - ): - items = items_callback(None, None) - self.assertEqual([item[0] for item in items], ["CPU"]) - - def test_remove_background_poll_requires_ben2_model(self): - context = SimpleNamespace( - object=SimpleNamespace( - type="EMPTY", - empty_display_type="IMAGE", - data=object(), - ) - ) - with patch.object( - self.addon.client, - "service_busy", - return_value=False, - ), patch.object( - self.addon.client, - "model_ready", - return_value=False, - ): - self.assertFalse(self.addon.RemoveImageBackground.poll(context)) - with patch.object( - self.addon.client, - "service_busy", - return_value=False, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ): - self.assertTrue(self.addon.RemoveImageBackground.poll(context)) - - def test_convert_to_depth_plane_poll_requires_da3_and_ben2_models(self): - context = SimpleNamespace( - object=SimpleNamespace( - type="EMPTY", - empty_display_type="IMAGE", - data=object(), - ) - ) - with patch.object( - self.addon.client, - "service_busy", - return_value=False, - ), patch.object( - self.addon.client, - "model_ready", - side_effect=lambda key: key in {"BASE", "BEN2_BASE"}, - ): - self.assertTrue( - self.addon.ConvertImageToDepthPlane.poll(context) - ) - with patch.object( - self.addon.client, - "service_busy", - return_value=False, - ), patch.object( - self.addon.client, - "model_ready", - return_value=False, - ): - self.assertTrue( - self.addon.ConvertImageToDepthPlane.poll(context) - ) - - def test_debug_panel_is_closed_by_default(self): - self.assertIn( - "DEFAULT_CLOSED", - self.addon.DebugPanel.bl_options, - ) - panels = ( - self.addon.Moge2ApiDebugPanel, - self.addon.Da3OutputDebugPanel, - self.addon.Ben2RemovalDebugPanel, - ) - for panel in panels: - self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname) - self.assertIn("DEFAULT_CLOSED", panel.bl_options) - self.assertEqual( - [panel.bl_label for panel in panels], - ["MoGe-2", "DA3", "BEN2"], - ) - - def test_service_panel_exposes_shared_service_state_and_controls(self): - self.assertIn(self.addon.AiServicePanel, self.addon.CLASSES) - self.assertIn(self.addon.ControlService, self.addon.CLASSES) - self.assertEqual(self.addon.AiServicePanel.bl_label, "AI Service") - - def test_models_is_an_independent_closed_panel(self): - self.assertIn(self.addon.ModelsPanel, self.addon.CLASSES) - self.assertNotIn("bl_parent_id", self.addon.ModelsPanel.__dict__) - self.assertIn("DEFAULT_CLOSED", self.addon.ModelsPanel.bl_options) - - def test_ai_operation_restarts_a_manually_stopped_service(self): - client = self.addon.client - controller = client.service_controller() - connection = SimpleNamespace(client=object()) - with patch.object(controller, "ensure", return_value=connection) as ensure: - result = client.ensure_service(timeout=0.1) - ensure.assert_called_once_with(0.1) - self.assertIs(result, connection) - - def test_blender_startup_does_not_start_ai_service(self): - client = self.addon.client - controller = client.service_controller() - with patch.object(client, "env_installed", return_value=True), patch.object( - controller, - "poll", - return_value=2.0, - ) as poll: - interval = client._service_timer() - poll.assert_called_once_with(available=True) - self.assertEqual(interval, 2.0) - - def test_ben2_debug_job_uses_its_image_and_model(self): - settings = SimpleNamespace( - background_input_path="input.png", - device="CUDA", - ) - operator = self.addon.operators.DebugBen2() - - with TemporaryDirectory() as directory: + def test_runtime_owns_environment_and_storage_paths(self): + runtime = self.addon.runtime + with tempfile.TemporaryDirectory() as directory: root = Path(directory) - input_path = root / "input.png" - input_path.touch() - settings.background_input_path = str(input_path) - paths = { - "log": root / "job.log", - "status": root / "job.json", - "result": root / "result", - } - with patch.object( - self.addon.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ) as model_ready, patch.object( - self.addon.client, - "debug_ben2", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as remove_background: - command = operator.create_request(settings) + with patch.object(runtime, "storage_root_factory", root): + self.assertEqual(runtime.storage_root(), root.resolve()) + self.assertEqual(runtime.environment_directory(), root / ".venv") + self.assertEqual(runtime.environment_status_path(), root / "environment-status.json") + self.assertEqual(runtime.environment_log_path(), root / "environment.log") + self.assertEqual(runtime.server_log_path(), root / "server.log") - self.assertIsInstance(command, self.addon.operators.JobRequest) - model_ready.assert_called_once_with("BEN2_BASE") - self.assertEqual(remove_background.call_args.args[1], input_path) - self.assertEqual(remove_background.call_args.args[2], "BEN2_BASE") - - def test_debug_jobs_use_the_scene_device(self): - settings = SimpleNamespace( - background_input_path="input.png", - device="CPU", - ) - operator = self.addon.operators.DebugBen2() - - with TemporaryDirectory() as directory: + def test_runtime_checks_its_environment_dictionary_hash(self): + runtime = self.addon.runtime + with tempfile.TemporaryDirectory() as directory: root = Path(directory) - input_path = root / "input.png" - input_path.touch() - settings.background_input_path = str(input_path) - paths = { - "log": root / "job.log", - "status": root / "job.json", - "result": root / "result", - } - with patch.object( - self.addon.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "debug_ben2", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as remove_background: - operator.create_request(settings) + python = root / ".venv" / ( + "Scripts/python.exe" if sys.platform == "win32" else "bin/python" + ) + python.parent.mkdir(parents=True) + python.touch() + (root / "manifest.json").write_text( + json.dumps({"environment_hash": runtime.environment_hash}), + encoding="utf-8", + ) + with patch.object(runtime, "storage_root_factory", root): + self.assertTrue(runtime.environment_ready()) - self.assertEqual( - remove_background.call_args.kwargs["device"], - "CPU", - ) - - def test_production_jobs_use_the_preferences_device(self): - settings = SimpleNamespace(device="CPU", video_frames=16) - operator = self.addon.operators.ProcessImageBackground() - - with TemporaryDirectory() as directory: + def test_runtime_builds_the_standard_installer_command(self): + runtime = self.addon.runtime + with tempfile.TemporaryDirectory() as directory: root = Path(directory) - input_path = root / "input.png" - 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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "production_device", - return_value="MPS", - ), patch.object( - self.addon.client, - "remove_background", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as remove_background: - operator.create_request(settings) + with patch.object(runtime, "storage_root_factory", root): + command = runtime.install_command() + self.assertEqual(command[0], sys.executable) + self.assertEqual(Path(command[1]).name, "environment.py") + self.assertEqual(Path(command[1]).parent.name, "installer") + self.assertIn("--storage-root", command) + self.assertIn("--config", command) + self.assertNotIn("--venv-dir", command) + self.assertNotIn("--models-dir", command) + self.assertEqual(command[command.index("--stages") + 1], "1") - self.assertEqual( - remove_background.call_args.kwargs["device"], - "MPS", - ) + def test_runtime_exposes_declared_environment_and_post_install(self): + runtime = self.addon.runtime + self.assertEqual(runtime.environment["python"], "3.12") + self.assertIn("numpy==2.4.2", runtime.environment["packages"]) + self.assertTrue(callable(runtime.post_install)) - def test_image_background_job_uses_operator_input(self): - settings = SimpleNamespace(device="AUTO") - operator = self.addon.operators.ProcessImageBackground() - - with TemporaryDirectory() as directory: - root = Path(directory) - input_path = root / "input.png" - 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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "remove_background", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as remove_background: - command = operator.create_request(settings) - - self.assertIsInstance(command, self.addon.operators.JobRequest) - 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.operators.ProcessImageBackground() - - 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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "model_ready", - return_value=True, - ), patch.object( - self.addon.client, - "remove_background", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as remove_background: - command = operator.create_request(settings) - - self.assertIsInstance(command, self.addon.operators.JobRequest) - 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.operators.ExtractOriginalFrames() - - 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.client, - "env_installed", - return_value=True, - ), patch.object( - self.addon.client, - "extract_frames", - return_value=self.addon.operators.JobRequest( - {"job_id": "test"} - ), - ) as extract_frames: - command = operator.create_request(settings) - - self.assertIsInstance(command, self.addon.operators.JobRequest) - 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): + def test_snapform_post_install_is_ordinary_code_using_runtime_request(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"}, - ), + fake_runtime = SimpleNamespace( + request=lambda job_type, parameters: calls.append( + (job_type, parameters) + ) ) - operator = self.addon.operators.SnapFormJobOperator() - job = SimpleNamespace( - result_directory=Path("result"), - source_object_name="Reference", - ) - - operator._apply_removed_background(job) - + self.addon.job_runtime.post_install(fake_runtime) self.assertEqual( calls, [ - ( - "apply", - { - "snapform_source_object_name": "Reference", - "snapform_output_path": str( - Path("result") / "foreground.png" - ), - }, - ), - ("undo", {"message": "Remove Background"}), + ("download-model", {"model": "MOGE2_VITS_NORMAL"}), + ("download-model", {"model": "BEN2_BASE"}), ], ) - def test_completed_image_background_job_applies_foreground_sequence(self): + def test_environment_controller_runs_arbitrary_post_install_callable(self): + runtime_module = importlib.import_module("snapform.blendjob.runtime") + runtime = self.addon.runtime 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"}, - ), + started = threading.Event() + release = threading.Event() + + def post_install(received): + calls.append(received) + started.set() + release.wait(1.0) + + controller = runtime_module.EnvironmentController(runtime) + controller.process = SimpleNamespace(poll=lambda: 0, returncode=0) + with tempfile.TemporaryDirectory() as directory, patch.object( + runtime, + "post_install", + post_install, + ), patch.object( + runtime, + "environment_status_path", + return_value=Path(directory) / "missing.json", + ): + self.assertEqual(controller.status("environment")["state"], "running") + self.assertTrue(started.wait(1.0)) + release.set() + controller.post_thread.join(1.0) + status = controller.status("environment") + + self.assertEqual(calls, [runtime]) + self.assertEqual(status["state"], "succeeded") + + def test_post_install_progress_resets_without_stage_fraction(self): + runtime_module = importlib.import_module("snapform.blendjob.runtime") + controller = runtime_module.EnvironmentController(self.addon.runtime) + + controller._update_post_progress( + {"job_id": "first", "job_type": "download-model", "progress": 1.0} ) - operator = self.addon.operators.SnapFormJobOperator() - 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"}), + first_stage = controller.post_status["stage"] + controller._update_post_progress( + {"job_id": "second", "job_type": "download-model", "progress": 0.1} ) - def test_completed_original_background_job_builds_with_one_undo(self): + self.assertGreater(controller.post_status["stage"], first_stage) + self.assertNotIn("stages", controller.post_status) + + def test_conversion_model_lists_default_to_moge2_vits(self): + properties = importlib.import_module("snapform.properties") + catalog = ( + {"key": "SMALL", "family": "da3", "enabled": True, + "label": "DA3 Small", "description": "DA3", "license": "Apache"}, + {"key": "MOGE2_VITS_NORMAL", "family": "moge2", "enabled": True, + "label": "MoGe-2 ViT-S", "description": "MoGe", "license": "MIT"}, + ) + with patch.object(properties, "model_catalog", return_value=catalog): + depth_items = properties.depth_plane_model_enum_items(None, None) + surface_items = properties.surface_model_enum_items(None, None) + + self.assertEqual(depth_items[0][0], "MOGE2_VITS_NORMAL") + self.assertEqual(surface_items[0][0], "MOGE2_VITS_NORMAL") + + def test_runtime_request_uses_request_response_shape(self): + result = object() + responses = [] + with patch.object( + self.addon.runtime.server, + "request", + return_value=result, + ) as request: + returned = self.addon.runtime.request( + "example", + {"value": 3}, + response=responses.append, + ) + self.assertIs(returned, result) + request.assert_called_once_with( + "example", + {"value": 3}, + response=responses.append, + progress=None, + cancel_check=None, + ) + + def test_server_request_returns_job_result_and_calls_response(self): + blendjob = importlib.import_module("snapform.blendjob") + controller = blendjob.ServerController.__new__(blendjob.ServerController) + statuses = iter( + ( + {"state": "running", "progress": 0.5}, + { + "state": "succeeded", + "progress": 1.0, + "result": {"value": 7}, + }, + ) + ) + controller.submit = lambda _job_type, _parameters: { + "job_id": "job", + "directory": ".", + } + controller.status = lambda _job_id: next(statuses) + controller.mark_job_complete = lambda _job_id: None + responses = [] + + result = controller.request( + "example", + {"input": 3}, + response=responses.append, + poll_interval=0.01, + ) + + self.assertEqual(result.value, {"value": 7}) + self.assertEqual(responses, [result]) + + def test_runtime_locates_the_package_exporting_server(self): + entrypoint = self.addon.runtime._server_entrypoint() + self.assertEqual(entrypoint.name, "__init__.py") + self.assertEqual(entrypoint.parent.name, "server") + + def test_runtime_registers_system_operators_and_status_bar(self): + self.addon.register() + system_classes = self.addon.runtime.operator_classes() + + self.assertTrue(set(system_classes) <= set(self.registered)) + self.assertEqual(len(self.status_bar_draws), 1) + self.assertEqual(self.status_bar_draws[0]._owner, "snapform") + + self.addon.unregister() + self.assertEqual(self.registered, []) + self.assertEqual(self.status_bar_draws, []) + + def test_status_bar_reads_runtime_state(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"}, - ), + + class Layout: + def row(self, **_options): + return self + + def progress(self, **options): + calls.append(("progress", options)) + + def operator(self, identifier, **_options): + calls.append(("operator", identifier)) + + runtime = self.addon.runtime + runtime.active_job = object() + runtime.progress = 0.4 + runtime.message = "Working" + runtime._draw_status_bar(SimpleNamespace(layout=Layout()), None) + runtime.active_job = None + + self.assertEqual(calls[0][1]["factor"], 0.4) + self.assertEqual(calls[0][1]["text"], "Working") + self.assertEqual(calls[1], ("operator", "snapform.cancel_job")) + + def test_moge_depth_request_is_written_on_the_operator(self): + operator = self.addon.operators.GenerateMoge2Depth() + operator.input_path = __file__ + operator.source_object_name = "" + operator.model = "MOGE2_VITS_NORMAL" + operator.resolution_level = 6 + operator.remove_background = False + settings = SimpleNamespace(video_frames=12) + context = SimpleNamespace( + scene=SimpleNamespace(snapform_settings=settings) ) - operator = self.addon.operators.SnapFormJobOperator() - job = SimpleNamespace( - kind="generate_original_background", - result_directory=Path("result"), + + with patch.object( + self.addon.operators, "require_environment" + ), patch.object(self.addon.operators, "require_model"): + parameters = operator.request(context) + + self.assertEqual(parameters["input"], __file__) + self.assertEqual(parameters["model"], "MOGE2_VITS_NORMAL") + self.assertEqual(parameters["resolution_level"], 6) + self.assertNotIn("job_type", parameters) + + def test_moge_depth_response_calls_an_ordinary_function(self): + operator = self.addon.operators.GenerateMoge2Depth() + operator.source_object_name = "Reference" + operator.remove_background = False + operator.delete_source = True + operator.unshaded = False + result = object() + context = object() + + with patch.object( + self.addon.operators, "build_moge2_depth_response" + ) as build: + operator.response(context, result) + + build.assert_called_once_with( + context, + result, source_object_name="Reference", - subdivision_level=5, - delete_source=False, - displacement_via="MODIFIER", + remove_background=False, + delete_source=True, unshaded=False, ) - operator._build_original_depth_plane(job) - - self.assertEqual(calls[0][0], "build") - self.assertEqual( - calls[0][1]["snapform_background_path"], - str(Path("result") / "foreground.png"), + def test_entry_operator_dispatches_plain_properties(self): + operator = self.addon.ConvertToDepthPlane() + operator.model = "MOGE2_VITS_NORMAL" + operator.resolution_level = 7 + operator.keep_original = False + operator.unshaded = True + operator.remove_background = False + source_object = SimpleNamespace( + name="Reference", + data=SimpleNamespace(source="FILE"), ) - self.assertFalse(calls[0][1]["snapform_delete_source"]) - self.assertEqual( - calls[0][1]["snapform_depth_npz_path"], - "", + settings = SimpleNamespace(video_frames=1) + context = SimpleNamespace( + object=source_object, + scene=SimpleNamespace(snapform_settings=settings), ) - self.assertEqual( - calls[1], - ("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.operators.SnapFormJobOperator() - with TemporaryDirectory() as directory: - result_directory = Path(directory) - (result_directory / "frame.0001.jpg").touch() - (result_directory / "frame.0002.jpg").touch() - job = SimpleNamespace( - kind="extract_original_frames", - result_directory=result_directory, - source_object_name="Reference", - subdivision_level=5, - delete_source=False, - displacement_via="MODIFIER", - unshaded=False, + generate_moge2_depth=lambda **options: calls.append(options) + or {"RUNNING_MODAL"} ) - - 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_uses_model_job(self): - client = importlib.import_module("snapform.client") - request = importlib.import_module("snapform.client").download_model( - "BEN2_BASE" - ) - - self.assertEqual(request.payload["job_type"], "download-model") - self.assertEqual(request.payload["parameters"]["model"], "BEN2_BASE") - - def test_install_env_command_has_no_model_label_reference(self): - client = importlib.import_module("snapform.client") with patch.object( - client, - "uv_executable", - return_value="C:/uv/uv.exe", - ), patch.object( - client, - "machine_devices", - return_value=["cuda", "cpu"], - ), patch.object( - client, - "venv_directory", - return_value=Path(".venv"), - ), patch.object( - client, - "models_directory", - return_value=Path("models"), + self.addon.image_depth, + "prepare_image_input", + return_value=(Path(__file__), False), ): - command = client.install_env_command(Path("status.json")) + self.assertEqual(operator.execute(context), {"FINISHED"}) - self.assertNotIn("--label", command) - self.assertEqual(command[0], sys.executable) - self.assertIn("--uv", command) - self.assertIn("C:/uv/uv.exe", command) - self.assertNotIn("--torch-variant", command) - self.assertTrue( - any("install_env.py" in str(value) for value in command) - ) + self.assertEqual(calls[0]["resolution_level"], 7) + self.assertEqual(calls[0]["source_object_name"], "Reference") + self.assertFalse(any(key.startswith("snapform_") for key in calls[0])) - def test_install_env_command_does_not_select_legacy_runtime(self): - client = importlib.import_module("snapform.client") - with patch.object( - client, - "uv_executable", - return_value="C:/uv/uv.exe", - ), patch.object( - client, - "machine_devices", - return_value=["cpu"], - ), patch.object( - client, - "venv_directory", - return_value=Path(".venv"), - ), patch.object( - client, - "models_directory", - return_value=Path("models"), - ): - command = client.install_env_command(Path("status.json")) + def test_models_are_owned_by_the_server_resource(self): + self.assertFalse((SOURCE_ROOT / "snapform" / "models.py").exists()) + with tempfile.TemporaryDirectory() as directory: + manager = self.addon.server.ModelManager(Path(directory)) + snapshot = manager.snapshot() + self.assertIn("catalog", snapshot) + self.assertTrue(snapshot["catalog"]) - self.assertNotIn("--pypi-index", command) - self.assertNotIn("--torch-variant", command) - self.assertIn("C:/uv/uv.exe", command) - def test_completed_export_opens_job_directory(self): - calls = [] - self.fake_bpy.ops = SimpleNamespace( - wm=SimpleNamespace( - path_open=lambda **options: calls.append(options) or {"FINISHED"}, - ), - ) - with TemporaryDirectory() as directory: - self.addon.operators.open_job_directory(directory) - - self.assertEqual(len(calls), 1) - self.assertIn("filepath", calls[0]) - - def test_status_bar_draws_progress_and_cancel_button(self): - progress_calls = [] - operator_calls = [] - - class Layout: - ui_units_x = 0.0 - - 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)) - - context = SimpleNamespace( - scene=SimpleNamespace( - snapform_settings=SimpleNamespace( - status="Processing", - progress=0.25, - ), - ), - ) - - job_operator = importlib.import_module("snapform.blendjob.operator") - with patch.object(job_operator, "_active_job", SimpleNamespace()): - self.addon.draw_status_bar( - SimpleNamespace(layout=Layout()), - context, - ) - - self.assertEqual( - progress_calls, - [{"factor": 0.25, "type": "BAR", "text": "Processing"}], - ) - self.assertEqual( - operator_calls, - [("snapform.cancel_job", {"text": "", "icon": "X"})], - ) - - def test_ui_redraw_targets_status_bar_and_view_3d_areas(self): - redrawn = [] - status_regions = ( - SimpleNamespace( - tag_redraw=lambda: redrawn.append("status-region"), - ), - ) - areas = ( - SimpleNamespace( - type="STATUSBAR", - tag_redraw=lambda: redrawn.append("status"), - regions=status_regions, - ), - SimpleNamespace( - type="VIEW_3D", - tag_redraw=lambda: redrawn.append("view"), - regions=(), - ), - ) - context = SimpleNamespace( - window_manager=SimpleNamespace( - windows=( - SimpleNamespace(screen=SimpleNamespace(areas=areas)), - ), - ), - ) - - self.addon.status_bar.redraw_ui(context) - self.addon.status_bar.redraw_ui(context) - - self.assertEqual( - redrawn, - [ - "status", - "status-region", - "view", - "status", - "status-region", - "view", - ], - ) - - def test_force_flag_refreshes_status_text(self): - redrawn = [] - context = SimpleNamespace( - window_manager=SimpleNamespace( - windows=( - SimpleNamespace( - screen=SimpleNamespace( - areas=( - SimpleNamespace( - type="STATUSBAR", - tag_redraw=lambda: redrawn.append( - "status" - ), - regions=(), - ), - ), - ), - ), - ), - ), - workspace=SimpleNamespace( - status_text_set_internal=lambda _value: redrawn.append( - "status-text" - ), - ), - ) - - self.addon.status_bar.redraw_ui(context, force=True) - - self.assertEqual( - redrawn, - [ - "status", - "status-text", - ], - ) - - def test_force_flag_ignores_missing_workspace(self): - redrawn = [] - context = SimpleNamespace( - window_manager=SimpleNamespace( - windows=( - SimpleNamespace( - screen=SimpleNamespace( - areas=( - SimpleNamespace( - type="STATUSBAR", - tag_redraw=lambda: redrawn.append( - "status" - ), - regions=(), - ), - ), - ), - ), - ), - ), - ) - - self.addon.status_bar.redraw_ui(context, force=True) - - self.assertEqual(redrawn, ["status"]) +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_ben2_onnx.py b/tests/test_ben2_onnx.py index 6decedf..c1a7680 100644 --- a/tests/test_ben2_onnx.py +++ b/tests/test_ben2_onnx.py @@ -11,11 +11,11 @@ SOURCE_DIR = ( Path(__file__).resolve().parents[1] / "src" / "snapform" - / "ai_service" + / "server" ) sys.path.insert(0, str(SOURCE_DIR)) -import ben2_onnx +import onnx_ben2 class FakeSession: @@ -35,7 +35,7 @@ class FakeSession: class Ben2OnnxTest(unittest.TestCase): def test_preprocess_matches_model_layout(self): image = Image.new("RGB", (40, 20), (64, 128, 192)) - result = ben2_onnx.preprocess(image) + result = onnx_ben2.preprocess(image) self.assertEqual(result.shape, (1, 3, 1024, 1024)) self.assertEqual(result.dtype, np.float32) @@ -45,7 +45,7 @@ class Ben2OnnxTest(unittest.TestCase): session = FakeSession() image = Image.new("RGB", (40, 20), (10, 20, 30)) - result = ben2_onnx.infer(session, image) + result = onnx_ben2.infer(session, image) self.assertEqual(result.mode, "RGBA") self.assertEqual(result.size, image.size) diff --git a/tests/test_da3_onnx.py b/tests/test_da3_onnx.py index a668274..6a007c3 100644 --- a/tests/test_da3_onnx.py +++ b/tests/test_da3_onnx.py @@ -12,11 +12,11 @@ SOURCE_DIR = ( Path(__file__).resolve().parents[1] / "src" / "snapform" - / "ai_service" + / "server" ) sys.path.insert(0, str(SOURCE_DIR)) -import da3_onnx +import onnx_da3 class FakeSession: @@ -37,7 +37,7 @@ class FakeSession: class Da3OnnxTest(unittest.TestCase): def test_target_size_uses_patch_multiple(self): self.assertEqual( - da3_onnx.target_size((800, 600), 504, "upper_bound_resize"), + onnx_da3.target_size((800, 600), 504, "upper_bound_resize"), (504, 378), ) @@ -47,7 +47,7 @@ class Da3OnnxTest(unittest.TestCase): Image.new("RGB", (80, 60), (10, 20, 30)).save(image_path) session = FakeSession() - result = da3_onnx.infer( + result = onnx_da3.infer( session, [image_path], process_res=56, diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index d640c56..e5add61 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -1,7 +1,5 @@ -import importlib.util -import json import os -import tempfile +import sys import unittest from pathlib import Path from unittest.mock import patch @@ -10,14 +8,16 @@ from unittest.mock import patch PROJECT_ROOT = Path(__file__).resolve().parents[1] -def load_worker(): - specification = importlib.util.spec_from_file_location( - "snapform_download_worker", - PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py", - ) - module = importlib.util.module_from_spec(specification) - specification.loader.exec_module(module) - return module +def load_download_module(): + addon_directory = PROJECT_ROOT / "src" / "snapform" + sys.path.insert(0, str(addon_directory)) + try: + return __import__( + "server.model_download", + fromlist=["model_download"], + ) + finally: + sys.path.remove(str(addon_directory)) class FakeTqdm: @@ -36,52 +36,41 @@ class FakeTqdm: class DownloadProgressTest(unittest.TestCase): @classmethod def setUpClass(cls): - cls.worker = load_worker() + cls.download = load_download_module() + + def create_progress(self, description="Downloading bytes", **options): + reports = [] + progress_class = self.download.create_download_progress_class( + lambda value, message: reports.append((value, message)), + "owner/model", + tqdm_type=options.pop("tqdm_type", FakeTqdm), + clock=options.pop("clock", lambda: 1.0), + **options, + ) + progress = progress_class( + desc=description, + total=100, + initial=0, + unit="B", + ) + return progress, reports def test_reports_aggregated_download_bytes(self): - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdm, - clock=lambda: 1.0, - ) - progress = progress_class( - desc="Downloading bytes", - total=100, - initial=0, - unit="B", - ) + progress, reports = self.create_progress() - progress.update(25) + progress.update(25) - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertAlmostEqual(status["progress"], 0.275) - self.assertEqual(status["message"], "Downloading owner/model") + self.assertAlmostEqual(reports[-1][0], 0.275) + self.assertEqual(reports[-1][1], "Downloading owner/model") - def test_reports_byte_progress_for_aggregated_bar_description(self): - # huggingface_hub 1.x 用 "Downloading (incomplete total...)" 作为字节进度条描述 - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdm, - clock=lambda: 1.0, - ) - progress = progress_class( - desc="Downloading (incomplete total...)", - total=100, - initial=0, - unit="B", - ) + def test_reports_hugging_face_aggregated_bar(self): + progress, reports = self.create_progress( + "Downloading (incomplete total...)" + ) - progress.update(25) + progress.update(25) - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertAlmostEqual(status["progress"], 0.275) - self.assertEqual(status["message"], "Downloading owner/model") + self.assertAlmostEqual(reports[-1][0], 0.275) def test_reports_download_speed(self): class FakeTqdmWithRate(FakeTqdm): @@ -89,129 +78,48 @@ class DownloadProgressTest(unittest.TestCase): def format_dict(self): return {"rate": 14_000_000} - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdmWithRate, - clock=lambda: 1.0, - label="da3 small", - ) - progress = progress_class( - desc="Downloading bytes", - total=100, - initial=0, - unit="B", - ) + progress, reports = self.create_progress( + tqdm_type=FakeTqdmWithRate, + label="da3 small", + ) - progress.update(50) + progress.update(50) - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertEqual( - status["message"], - "Downloading da3 small 13.4 MB/s", - ) + self.assertEqual( + reports[-1][1], + "Downloading da3 small 13.4 MB/s", + ) - def test_does_not_report_estimated_time_remaining(self): + def test_progress_callback_has_no_eta_or_state_protocol(self): times = iter((0.0, 10.0)) - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdm, - clock=lambda: next(times), - ) - progress = progress_class( - desc="Downloading bytes", - total=100, - initial=0, - unit="B", - ) + progress, reports = self.create_progress(clock=lambda: next(times)) - progress.update(50) + progress.update(50) - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertNotIn("eta_seconds", status) - - def test_staged_download_keeps_model_label(self): - class FakeTqdmWithRate(FakeTqdm): - @property - def format_dict(self): - return {"rate": 14_000_000} - - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdmWithRate, - clock=lambda: 1.0, - label="da3 small", - stage=6, - stages=8, - stage_label="da3 small", - ) - progress = progress_class( - desc="Downloading bytes", - total=100, - initial=0, - unit="B", - ) - - progress.update(50) - - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertEqual( - status["message"], - "Downloading da3 small 13.4 MB/s", - ) - self.assertEqual(status["stage"], 6) - self.assertEqual(status["stages"], 8) - self.assertEqual(status["stage_label"], "da3 small") + self.assertEqual(len(reports[-1]), 2) def test_ignores_unrelated_progress_bars(self): - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - progress_class = self.worker.create_download_progress_class( - status_path, - "owner/model", - tqdm_type=FakeTqdm, - clock=lambda: 1.0, - ) - progress = progress_class( - desc="Reconstructing", - total=100, - initial=0, - unit="B", - ) + progress, reports = self.create_progress("Reconstructing") - progress.update(25) + progress.update(25) - self.assertFalse(status_path.exists()) + self.assertEqual(reports, []) def test_download_endpoints_uses_direct_endpoint(self): with patch.dict(os.environ, {"HF_ENDPOINT": ""}): - endpoints = self.worker.download_endpoints() + endpoints = self.download.download_endpoints() - self.assertEqual( - endpoints, - ["https://huggingface.co"], - ) + self.assertEqual(endpoints, ["https://huggingface.co"]) def test_download_endpoints_respects_environment_override(self): with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example"}): - endpoints = self.worker.download_endpoints() + endpoints = self.download.download_endpoints() - self.assertEqual( - endpoints, - ["https://hf.example"], - ) + self.assertEqual(endpoints, ["https://hf.example"]) def test_download_endpoints_strips_trailing_slash(self): with patch.dict(os.environ, {"HF_ENDPOINT": "https://hf.example/"}): - endpoints = self.worker.download_endpoints() + endpoints = self.download.download_endpoints() self.assertEqual(endpoints, ["https://hf.example"]) diff --git a/tests/test_install_env.py b/tests/test_install_env.py index 5ea2c12..07c4695 100644 --- a/tests/test_install_env.py +++ b/tests/test_install_env.py @@ -1,1347 +1,120 @@ import importlib.util -import io -import hashlib import json -import os -import shutil -import subprocess -import tarfile +import sys import tempfile -import time import unittest -import zipfile from pathlib import Path -from types import SimpleNamespace -from unittest.mock import Mock, patch +from unittest.mock import patch PROJECT_ROOT = Path(__file__).resolve().parents[1] -def load_install_env(): +def load_environment(): + path = PROJECT_ROOT / "src" / "snapform" / "blendjob" / "environment.py" specification = importlib.util.spec_from_file_location( - "snapform_install_env", - PROJECT_ROOT / "src" / "snapform" / "install_env.py", + "blendjob_environment_test", + path, ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) return module -class InstallEnvTest(unittest.TestCase): - def create_bundled_vendor(self, root): - vendor = root / "vendor" - vendor.mkdir() - (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") - (vendor / "ben2.tar.gz").write_bytes(b"ben2") - (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") - (vendor / "moge.tar.gz").write_bytes(b"moge") - return vendor - - def test_installs_onnx_runtime_without_training_frameworks(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ) as mocked_run, patch.object( - bootstrap, - "run_install", - ) as mocked_run_install, patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - commands = [call.args[0] for call in mocked_run.call_args_list] - commands += [ - call.args[4] for call in mocked_run_install.call_args_list - ] - service_command = next( - command - for command in commands - if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command - ) - self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", service_command) - self.assertIn( - f"onnxruntime-directml=={bootstrap.ONNXRUNTIME_VERSION}", - service_command, - ) - for forbidden in ("torch", "torchvision", "xformers", "ben2 @", "moge @"): - self.assertFalse( - any(forbidden in str(value) for value in service_command), - forbidden, - ) - check_command = bootstrap.verify_command( - bootstrap.venv_python(venv_dir) - ) - self.assertFalse(any("torch" in str(value) for value in check_command)) - self.assertTrue( - any("import fastapi, uvicorn" in str(value) for value in check_command) - ) - self.assertTrue( - any("onnxruntime" in str(value) for value in check_command) - ) - self.assertFalse( - any("import gsplat" in str(value) for value in check_command) - ) - self.assertTrue( - any( - "SNAPFORM_DEVICES=" in str(value) - for value in check_command - ) - ) - default_downloads = [ - command - for command in commands - if any( - repository in command - for repository in ( - "Ruicheng/moge-2-vits-normal-onnx", - "onnx-community/BEN2-ONNX", - ) - ) - ] - self.assertEqual(len(default_downloads), 2) - self.assertTrue(any("Ruicheng/moge-2-vits-normal-onnx" in map(str, command) for command in default_downloads)) - ben2_download = next( - command - for command in default_downloads - if "onnx-community/BEN2-ONNX" in command - ) - patterns = [ - ben2_download[index + 1] - for index, value in enumerate(ben2_download) - if value == "--allow-pattern" - ] - self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) - self.assertIn("--stage", ben2_download) - self.assertIn("--stage-label", ben2_download) - self.assertEqual( - ben2_download[ben2_download.index("--label") + 1], - "BEN2 Base", - ) - - def test_darwin_install_uses_default_onnx_runtime(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object( - bootstrap.sys, - "platform", - "darwin", - ), patch.object( - bootstrap.platform, - "machine", - return_value="arm64", - ), patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_install", - ) as mocked_run_install, patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - commands = [ - call.args[4] for call in mocked_run_install.call_args_list - ] - runtime_command = next( - command - for command in commands - if any( - f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}" in str(value) - for value in command - ) - ) - self.assertNotIn("--index-url", runtime_command) - self.assertFalse( - any( - "xformers" in str(value) - for command in commands - for value in command - ) - ) - - def test_intel_mac_uses_portable_onnx_runtime(self): - bootstrap = load_install_env() - with patch.object( - bootstrap.sys, - "platform", - "darwin", - ): - stages = bootstrap.install_stages( - SimpleNamespace( - models_dir="models", - status="status.json", - ), - ".venv", - "python", - "uv", - ) - runtime = stages[1][1] - self.assertIn( - f"onnxruntime=={bootstrap.ONNXRUNTIME_VERSION}", - runtime, - ) - - def test_run_stage_writes_stage_progress(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - with patch.object(bootstrap, "run") as mocked_run: - bootstrap.run_stage( - str(status_path), - 3, - 8, - "pytorch", - ["uv", "pip", "install", "torch"], - None, - "plain", - ) - - self.assertEqual( - mocked_run.call_args.args[0], - ["uv", "pip", "install", "torch"], - ) - status = json.loads(status_path.read_text(encoding="utf-8")) - self.assertEqual(status["progress"], 1.0) - self.assertEqual(status["stage"], 3) - self.assertEqual(status["stages"], 8) - self.assertEqual(status["stage_label"], "pytorch") - - def test_run_stage_writes_progress_range(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - with patch.object(bootstrap, "run"), patch.object( - bootstrap, - "write_status", - ) as mocked_write: - bootstrap.run_stage( - str(status_path), - 3, - 8, - "pytorch", - ["uv", "pip", "install", "torch"], - None, - "plain", - start=0.2, - end=0.4, - ) - - progresses = [ - call.args[1] for call in mocked_write.call_args_list - ] - self.assertEqual(progresses[0], 0.2) - self.assertEqual(progresses[-1], 0.4) - - def test_run_stage_passes_progress_range_to_download_command(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - command = [ - "python", - "worker.py", - "download", - "--repository", - "depth-anything/DA3-SMALL", - ] - with patch.object(bootstrap, "run") as mocked_run: - bootstrap.run_stage( - str(status_path), - 6, - 8, - "da3 small", - command, - None, - "download", - start=0.9, - end=0.9333, - ) - - called_command = mocked_run.call_args.args[0] - self.assertIn("--progress-start", called_command) - self.assertIn("--progress-end", called_command) - self.assertEqual( - called_command[called_command.index("--progress-start") + 1], - "0.9", - ) - self.assertEqual( - called_command[called_command.index("--progress-end") + 1], - "0.9333", - ) - self.assertEqual(called_command[-1], "da3 small") - - def test_install_resets_progress_for_each_stage(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - stage_ranges = [] - - def fake_run_stage( - status_path, - stage, - stages, - label, - command, - environment, - kind, - start=0.0, - end=1.0, - ): - stage_ranges.append((stage, start, end)) - bootstrap.write_status( - status_path, - end, - "", - stage=stage, - stages=stages, - stage_label=label, - ) - return [] - - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run_stage", - side_effect=fake_run_stage, - ), patch.object(bootstrap, "run"), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - self.assertEqual(stage_ranges[:2], [(1, 0.0, 0.05), (1, 0.05, 1.0)]) - self.assertEqual(stage_ranges[2:], [(2, 0.0, 1.0), (3, 0.0, 1.0)]) - status = json.loads( - (root / "status.json").read_text(encoding="utf-8") - ) - self.assertEqual(status["progress"], 1.0) - - def test_run_install_tracks_package_download_events(self): - bootstrap = load_install_env() - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - fake_process = FakeProcess( - [ - "Resolved 3 packages in 1.0s", - "DEBUG Adding transitive dependency for torch==2.9.1: filelock*", - "Prepared 3 packages in 2.0s", - "Installed 3 packages in 0.5s", - "+ torch==2.9.1", - "+ numpy==1.26.4", - ] - ) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ), patch.object(bootstrap, "write_status") as mocked_write: - installed = bootstrap.run_install( - "status.json", - 2, - 8, - "pytorch deps", - ["uv", "pip", "install", "torch"], - ) - - messages = [call.args[2] for call in mocked_write.call_args_list] - self.assertEqual( - messages, - [ - "Resolving 3 packages", - "Downloading packages 3/3", - "Installing packages 3/3", - ], - ) - 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) - self.assertEqual(first.kwargs["stage_label"], "pytorch deps") - self.assertEqual(installed, ["torch==2.9.1", "numpy==1.26.4"]) - - def test_run_install_smooths_progress_between_events(self): - bootstrap = load_install_env() - - class SlowLines: - def __init__(self, lines, delay): - self.lines = list(lines) - self.delay = delay - - def __iter__(self): - for line in self.lines: - time.sleep(self.delay) - yield line - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - fake_process = FakeProcess( - SlowLines( - [ - "Resolved 3 packages in 1.0s", - "Prepared 3 packages in 2.0s", - "Installed 3 packages in 0.5s", - ], - 0.3, - ) - ) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ), patch.object(bootstrap, "write_status") as mocked_write: - bootstrap.run_install( - "status.json", - 2, - 8, - "pytorch deps", - ["uv", "pip", "install", "torch"], - ) - - progresses = [call.args[1] for call in mocked_write.call_args_list] - self.assertGreater(progresses[-1], 0.02) - self.assertTrue( - all( - progresses[i] <= progresses[i + 1] + 1e-9 - for i in range(len(progresses) - 1) - ) - ) - messages = [call.args[2] for call in mocked_write.call_args_list] - self.assertIn("Resolving 3 packages", messages) - self.assertIn("Downloading packages 3/3", messages) - - def test_run_install_maps_progress_into_stage_range(self): - bootstrap = load_install_env() - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - fake_process = FakeProcess( - [ - "Resolved 3 packages in 1.0s", - "Prepared 3 packages in 2.0s", - "Installed 3 packages in 0.5s", - ] - ) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ), patch.object(bootstrap, "write_status") as mocked_write: - bootstrap.run_install( - "status.json", - 2, - 8, - "pytorch deps", - ["uv", "pip", "install", "torch"], - start=0.2, - end=0.4, - ) - - first = mocked_write.call_args_list[0] - self.assertAlmostEqual(first.args[1], 0.204) - self.assertEqual(first.kwargs["stage"], 2) - self.assertEqual(first.kwargs["stages"], 8) - self.assertEqual(first.kwargs["stage_label"], "pytorch deps") - - def test_run_install_heartbeat_reports_package_count(self): - bootstrap = load_install_env() - - class SlowLines: - def __iter__(self): - yield "Resolved 3 packages in 1.0s\n" - time.sleep(0.6) - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - clock_state = {"now": 0.0} - - def fake_clock(): - clock_state["now"] += 6.0 - return clock_state["now"] - - fake_process = FakeProcess(SlowLines()) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ), patch.object(bootstrap, "write_status") as mocked_write: - bootstrap._run_install_once( - "status.json", - 2, - 2, - "ONNX Runtime", - ["uv", "pip", "install", "onnxruntime"], - None, - 0.10, - 0.85, - fake_clock, - ) - - heartbeats = [ - call - for call in mocked_write.call_args_list - if call.args[2].startswith("Downloading") - ] - self.assertTrue(heartbeats) - self.assertTrue( - any("3 packages" in call.args[2] for call in heartbeats) - ) - for call in heartbeats: - self.assertNotIn("elapsed", call.args[2]) - self.assertNotIn("eta_seconds", call.kwargs) - - def test_run_install_retries_across_indexes(self): - bootstrap = load_install_env() - 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[0].args[5], - {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[0]}, - ) - self.assertEqual( - mocked_install.call_args_list[1].args[5], - {"UV_INDEX_URL": bootstrap.PYPI_MIRRORS[1]}, - ) - - def test_run_install_uses_only_configured_index(self): - bootstrap = load_install_env() - previous = bootstrap._resolved_pypi_index - bootstrap._resolved_pypi_index = ( - "https://pypi.tuna.tsinghua.edu.cn/simple" - ) - try: - with patch.object( - bootstrap, - "_run_install_once", - return_value=["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, 1) - self.assertEqual( - mocked_install.call_args.args[5], - {"UV_INDEX_URL": "https://pypi.tuna.tsinghua.edu.cn/simple"}, - ) - - def test_resolve_indexes_uses_configured_pypi_index(self): - bootstrap = load_install_env() - bootstrap.resolve_indexes( - SimpleNamespace(pypi_index="https://pypi.tuna.tsinghua.edu.cn/simple") - ) - self.assertEqual( - bootstrap._resolved_pypi_index, - "https://pypi.tuna.tsinghua.edu.cn/simple", - ) - - def test_resolve_indexes_uses_environment_override(self): - bootstrap = load_install_env() - with patch.dict( - os.environ, - {"UV_INDEX_URL": "https://pypi.example/simple"}, - ): - bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) - self.assertEqual( - bootstrap._resolved_pypi_index, - "https://pypi.example/simple", - ) - - def test_resolve_indexes_falls_back_to_pip_index_url(self): - bootstrap = load_install_env() - with patch.dict( - os.environ, - {"PIP_INDEX_URL": "https://pypi.example/simple"}, - ): - bootstrap.resolve_indexes(SimpleNamespace(pypi_index=None)) - self.assertEqual( - bootstrap._resolved_pypi_index, - "https://pypi.example/simple", - ) - - def test_child_environment_injects_resolved_pypi_index(self): - bootstrap = load_install_env() - previous_pypi = bootstrap._resolved_pypi_index - bootstrap._resolved_pypi_index = ( - "https://pypi.tuna.tsinghua.edu.cn/simple" - ) - try: - env = bootstrap.child_environment() - finally: - bootstrap._resolved_pypi_index = previous_pypi - - self.assertEqual( - env["UV_INDEX_URL"], - "https://pypi.tuna.tsinghua.edu.cn/simple", - ) - self.assertEqual( - env["PYTHONIOENCODING"], - "utf-8:backslashreplace", - ) - self.assertEqual(env["PYTHONUTF8"], "1") - - def test_standard_streams_are_reconfigured_for_unicode_logs(self): - bootstrap = load_install_env() - stdout = Mock() - stderr = Mock() - - with patch.object(bootstrap.sys, "stdout", stdout), patch.object( - bootstrap.sys, - "stderr", - stderr, - ): - bootstrap.configure_standard_streams() - - stdout.reconfigure.assert_called_once_with( - encoding="utf-8", - errors="backslashreplace", - ) - stderr.reconfigure.assert_called_once_with( - encoding="utf-8", - errors="backslashreplace", - ) - - def test_ready_json_records_installed_packages(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_stage", - return_value=["torch==2.9.1+cu128", "numpy==1.26.4"], - ), patch.object( - bootstrap, - "run_verify", - return_value={"cuda": True, "mps": False}, - ): - bootstrap.install(args) - - ready = json.loads( - (root / "manifest.json").read_text( - encoding="utf-8" - ) - ) - self.assertEqual( - ready["packages"], - ["numpy==1.26.4", "torch==2.9.1+cu128"], - ) - self.assertEqual(ready["devices"], ["directml", "cpu"]) - self.assertEqual( - ready["manifest_revision"], - bootstrap.MANIFEST_REVISION, - ) - - def test_broken_env_is_cleaned_before_reinstall(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - (venv_dir / "partial-python.exe").write_text("junk") - args = SimpleNamespace( - uv="uv", - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_install", - ), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - self.assertFalse((venv_dir / "partial-python.exe").exists()) - self.assertTrue( - (root / "manifest.json").is_file() - ) - - def test_installed_env_is_not_cleaned(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir(parents=True) - self.create_bundled_vendor(root) - (venv_dir / "Scripts").mkdir() - (venv_dir / "Scripts" / "python.exe").write_text("python") - (root / "manifest.json").write_text( - json.dumps( - {"manifest_revision": bootstrap.MANIFEST_REVISION} - ), - encoding="utf-8", - ) - (venv_dir / "marker.txt").write_text("keep") - args = SimpleNamespace( - uv="uv", - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_install", - ), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - self.assertTrue((venv_dir / "marker.txt").is_file()) - - def test_model_download_failure_keeps_env_installed(self): - bootstrap = load_install_env() - commands = [] - - def fake_run(command, environment=None): - commands.append(command) - if any("moge-2-vits-normal" in str(value) for value in command): - raise subprocess.CalledProcessError(1, command) - - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - args = SimpleNamespace( - uv="uv", - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - side_effect=fake_run, - ), patch.object(bootstrap, "run_install"), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - self.assertTrue( - (root / "manifest.json").is_file() - ) - status = json.loads( - (root / "status.json").read_text(encoding="utf-8") - ) - self.assertEqual(status["progress"], 1.0) - self.assertIn("Environment is installed", status["message"]) - self.assertTrue( - any( - "onnx-community/BEN2-ONNX" in map(str, command) - for command in commands - ) - ) - - def test_install_fails_when_all_uv_sources_are_unavailable(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - uv_path = root / "uv" / "uv.exe" - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_install", - ), patch.object( - bootstrap.shutil, - "which", - return_value=None, - ), patch.object( - bootstrap.shutil, - "rmtree", - ), patch.object( - bootstrap, - "download_uv_archive", - side_effect=RuntimeError("all sources unavailable"), - ): - with self.assertRaisesRegex( - RuntimeError, - "all sources unavailable", - ): - bootstrap.install(args) - - def test_install_downloads_and_extracts_uv(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - target, archive_name = bootstrap.uv_download_target() - archive = root / f"uv-{target}.{archive_name}" - if archive_name.endswith("zip"): - with zipfile.ZipFile(archive, "w") as bundle: - bundle.writestr("uv/uv.exe", b"uv-binary") - 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), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="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 [] - - def fake_download(destination, *_args): - shutil.copy2(archive, destination) - return destination - - with patch.object( - bootstrap, - "resolve_indexes", - return_value=None, - ), patch.object( - bootstrap, - "download_uv_archive", - side_effect=fake_download, - ), patch.object( - bootstrap, - "run", - side_effect=fake_run, - ), patch.object( - bootstrap, - "run_install", - side_effect=fake_run_install, - ), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - self.assertEqual(uv_path.read_bytes(), b"uv-binary") - all_commands = run_commands + install_commands - runtime_command = next( - command - for command in install_commands - if any("onnxruntime" in str(value) for value in command) - ) - self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) - self.assertFalse(any("torch" in str(value) for value in runtime_command)) - status = json.loads(Path(args.status).read_text(encoding="utf-8")) - self.assertEqual(status["stage"], 4) - self.assertEqual(status["stages"], 4) - - def test_install_reuses_existing_uv(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv.exe" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - with patch.object(bootstrap, "resolve_indexes", return_value=None), patch.object( - bootstrap, - "run", - ), patch.object( - bootstrap, - "run_install", - ), patch.object( - bootstrap, - "run_verify", - return_value={}, - ): - bootstrap.install(args) - - def test_extract_uv_unpacks_standalone_binary(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - archive_path = root / "uv.zip" - with zipfile.ZipFile(archive_path, "w") as bundle: - bundle.writestr("uv-x86_64-pc-windows-msvc/uv.exe", "binary") - - uv_path = root / "uv.exe" - bootstrap.extract_uv(archive_path, uv_path) - - self.assertEqual(uv_path.read_bytes(), b"binary") - - def test_uv_download_falls_back_to_omoolab_and_verifies_checksum(self): - bootstrap = load_install_env() - content = io.BytesIO() - with zipfile.ZipFile(content, "w") as archive: - archive.writestr("uv/uv.exe", b"uv-binary") - payload = content.getvalue() - - class Response(io.BytesIO): - headers = {"Content-Length": str(len(payload))} - - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - destination = root / "uv.tmp" - status = root / "status.json" - urls = [] - clock_state = {"now": 0.0} - - def clock(): - clock_state["now"] += 0.25 - return clock_state["now"] - - def open_url(request): - urls.append(request.full_url) - if len(urls) == 1: - raise OSError("official unavailable") - return Response(payload) - - with patch.object( - bootstrap, - "uv_download_target", - return_value=("test-target", "zip"), - ), patch.dict( - bootstrap.UV_ARTIFACTS, - { - "test-target": ( - "uv-test.zip", - hashlib.sha256(payload).hexdigest(), - ) - }, - ), patch.dict( - bootstrap._uv_download_openers, - { - bootstrap.UV_OFFICIAL_BASE_URL: SimpleNamespace( - open=open_url - ), - bootstrap.UV_MIRROR_BASE_URL: SimpleNamespace( - open=open_url - ), - }, - ), patch.object( - bootstrap, - "write_status", - ) as write_status: - result = bootstrap.download_uv_archive( - destination, - status, - 0.0, - 1.0, - {}, - clock=clock, - ) - - self.assertEqual(result.read_bytes(), payload) - self.assertTrue(urls[0].startswith(bootstrap.UV_OFFICIAL_BASE_URL)) - self.assertTrue(urls[1].startswith(bootstrap.UV_MIRROR_BASE_URL)) - self.assertEqual(bootstrap.UV_OFFICIAL_CONNECT_TIMEOUT, 3.0) - self.assertEqual(bootstrap.UV_MIRROR_CONNECT_TIMEOUT, 15.0) - messages = [call.args[2] for call in write_status.call_args_list] - self.assertTrue(any("uv (mirror)" in value for value in messages)) - self.assertTrue(any("B/s" in value for value in messages)) - self.assertFalse(any("left" in value for value in messages)) - - def test_install_stages_use_only_runtime_packages(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - - stages = bootstrap.install_stages( - args, - venv_dir, - root / "python", - uv_path, - ) - - runtime_command = next( - command - for _label, command, _environment, kind in stages - if kind == "install" - and any("onnxruntime" in str(value) for value in command) - ) - self.assertFalse(any(" @ file:" in str(value) for value in runtime_command)) - self.assertFalse(any("torch" in str(value) for value in runtime_command)) - - def test_install_stages_download_only_default_models(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - vendor = root / "vendor" - vendor.mkdir() - (vendor / "depth-anything-3.tar.gz").write_bytes(b"da3") - (vendor / "ben2.tar.gz").write_bytes(b"ben2") - (vendor / "utils3d.tar.gz").write_bytes(b"utils3d") - (vendor / "moge.tar.gz").write_bytes(b"moge") - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) - - stages = bootstrap.install_stages( - args, - venv_dir, - root / "python", - uv_path, - ) - - downloads = { - label: command - for label, command, _environment, kind in stages - if kind == "download" +class EnvironmentTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.environment = load_environment() + + def config(self): + return { + "python": "3.12", + "packages": ["numpy==2.4.2"], + "platform_packages": { + "windows": ["onnxruntime-directml==1.24.4"], + "default": ["onnxruntime==1.24.4"], + }, } - self.assertIn( - "moge-2-vits-normal-onnx", - " ".join(map(str, downloads["MoGe-2 ViT-S Normal"])), - ) - self.assertIn("BEN2-ONNX", " ".join(map(str, downloads["BEN2 Base"]))) + + def test_normalizes_public_dictionary(self): self.assertEqual( - set(downloads), - {"MoGe-2 ViT-S Normal", "BEN2 Base"}, + self.environment.normalized_environment(self.config()), + self.config(), ) - def test_install_does_not_verify_env(self): - bootstrap = load_install_env() + def test_requires_python_and_list_packages(self): + with self.assertRaises(ValueError): + self.environment.normalized_environment({}) + with self.assertRaises(TypeError): + self.environment.normalized_environment( + {"python": "3.12", "packages": "numpy"} + ) + + def test_environment_digest_tracks_configuration(self): + first = self.environment.environment_digest(self.config()) + reordered = { + "platform_packages": self.config()["platform_packages"], + "packages": ["numpy==2.4.2"], + "python": "3.12", + } + self.assertEqual(first, self.environment.environment_digest(reordered)) + reordered["packages"].append("pillow==12.1.1") + self.assertNotEqual(first, self.environment.environment_digest(reordered)) + + def test_selects_platform_packages_and_server_packages(self): + with patch.object(self.environment.sys, "platform", "win32"): + packages = self.environment.resolved_packages(self.config()) + self.assertIn("fastapi==0.139.2", packages) + self.assertIn("uvicorn==0.51.0", packages) + self.assertIn("numpy==2.4.2", packages) + self.assertIn("onnxruntime-directml==1.24.4", packages) + self.assertNotIn("onnxruntime==1.24.4", packages) + + def test_uses_default_platform_packages(self): + with patch.object(self.environment.sys, "platform", "linux"): + packages = self.environment.resolved_packages(self.config()) + self.assertIn("onnxruntime==1.24.4", packages) + + def test_install_writes_manifest_from_environment_hash(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) - venv_dir = root / ".venv" - venv_dir.mkdir() - self.create_bundled_vendor(root) - uv_path = root / "uv" - uv_path.write_bytes(b"uv") - args = SimpleNamespace( - uv=str(uv_path), - venv_dir=str(venv_dir), - models_dir=str(root / "models"), - torch_variant="cu128", - da3_commit="test-commit", - ben2_commit="test-ben2-commit", - status=str(root / "status.json"), - ) + uv = root / "tools" / "uv" + uv.parent.mkdir() + uv.touch() + status = root / "environment-status.json" with patch.object( - bootstrap, - "resolve_indexes", - return_value=None, - ), patch.object(bootstrap, "run"), patch.object( - bootstrap, - "run_install", + self.environment, + "install_uv", + return_value=uv, ), patch.object( - bootstrap, - "run_verify", - ) as mocked_verify: - bootstrap.install(args) + self.environment, + "run_command", + ), patch.object( + self.environment, + "install_packages", + ) as install_packages: + self.environment.install(root, self.config(), status, stages=3) - mocked_verify.assert_not_called() + manifest = json.loads((root / "manifest.json").read_text("utf-8")) + final_status = json.loads(status.read_text("utf-8")) + self.assertEqual( + manifest["environment_hash"], + self.environment.environment_digest(self.config()), + ) + self.assertIn("numpy==2.4.2", manifest["packages"]) + self.assertEqual(final_status["stage"], 1) + self.assertEqual(final_status["stages"], 3) + self.assertEqual(final_status["progress"], 1.0) + self.assertTrue(install_packages.called) - def test_run_verify_returns_detected_devices(self): - bootstrap = load_install_env() - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - fake_process = FakeProcess( - ['SNAPFORM_DEVICES={"cuda": false, "mps": true}'] + def test_environment_python_is_inside_runtime_storage(self): + path = self.environment.environment_python("storage") + expected = ( + Path("storage/.venv/Scripts/python.exe") + if sys.platform == "win32" + else Path("storage/.venv/bin/python") ) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ): - report = bootstrap.run_verify( - "python", - "status.json", - clock=lambda: 0.0, - ) - - self.assertFalse(report["cuda"]) - self.assertTrue(report["mps"]) - - def test_run_verify_writes_heartbeat_progress(self): - bootstrap = load_install_env() - - class FakeProcess: - def __init__(self, lines): - self.stdout = lines - - def wait(self): - return 0 - - clock_state = {"now": 0.0} - - def fake_clock(): - clock_state["now"] += 6.0 - return clock_state["now"] - - def delayed_empty_lines(): - time.sleep(0.6) - if False: - yield - - fake_process = FakeProcess(delayed_empty_lines()) - with patch.object( - bootstrap.subprocess, - "Popen", - return_value=fake_process, - ), patch.object(bootstrap, "write_status") as mocked_write: - report = bootstrap.run_verify( - "python", - "status.json", - clock=fake_clock, - ) - - messages = [call.args[2] for call in mocked_write.call_args_list] - self.assertTrue( - any( - "Checking Environment" in message - for message in messages - ) - ) - self.assertEqual(report, {}) - - def test_install_uv_extracts_downloaded_archive(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - uv_path = root / "uv" - target, archive_name = bootstrap.uv_download_target() - archive = root / f"uv-{target}.{archive_name}" - if archive_name.endswith("zip"): - with zipfile.ZipFile(archive, "w") as bundle: - bundle.writestr("uv/uv.exe", b"uv-binary") - 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)) - def fake_download(destination, *_args): - shutil.copy2(archive, destination) - return destination - - with patch.object( - bootstrap, - "download_uv_archive", - side_effect=fake_download, - ): - bootstrap.install_uv(uv_path, str(root / "status.json")) - - self.assertEqual(uv_path.read_bytes(), b"uv-binary") - - def test_install_uv_raises_when_download_fails(self): - bootstrap = load_install_env() - with tempfile.TemporaryDirectory() as directory: - root = Path(directory) - uv_path = root / "uv" - with patch.object( - bootstrap, - "download_uv_archive", - side_effect=RuntimeError("all uv sources failed"), - ): - with self.assertRaisesRegex( - RuntimeError, - "all uv sources failed", - ): - bootstrap.install_uv(uv_path, str(root / "status.json")) + self.assertEqual(path, expected) if __name__ == "__main__": diff --git a/tests/test_models.py b/tests/test_models.py index d865f84..d64de33 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -21,7 +21,7 @@ class ModelTest(unittest.TestCase): def setUpClass(cls): cls.models = load_module( "snapform_models", - "src/snapform/models.py", + "src/snapform/server/model_catalog.py", ) def test_base_model_is_default_friendly_apache_model(self): diff --git a/tests/test_moge_onnx.py b/tests/test_moge_onnx.py index aa1f529..c7a2537 100644 --- a/tests/test_moge_onnx.py +++ b/tests/test_moge_onnx.py @@ -12,8 +12,8 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_module(): specification = importlib.util.spec_from_file_location( - "snapform_moge_onnx", - PROJECT_ROOT / "src" / "snapform" / "ai_service" / "moge_onnx.py", + "snapform_onnx_moge2", + PROJECT_ROOT / "src" / "snapform" / "server" / "onnx_moge2.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) diff --git a/tests/test_pack.py b/tests/test_pack.py index 1d57756..00fc470 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -58,8 +58,23 @@ class PackTest(unittest.TestCase): "assets/O_SnapForm.blend", names, ) - self.assertIn("ai_service/jobs.py", names) - self.assertIn("ai_service/server.py", names) + self.assertIn("server/media.py", names) + self.assertIn("server/geometry.py", names) + self.assertIn("server/background.py", names) + self.assertIn("server/moge2.py", names) + self.assertIn("server/da3.py", names) + self.assertNotIn("server/processing.py", names) + self.assertNotIn("server/progress.py", names) + self.assertIn("server/model_catalog.py", names) + self.assertIn("server/model_download.py", names) + self.assertIn("server/onnx_ben2.py", names) + self.assertIn("server/onnx_da3.py", names) + self.assertIn("server/onnx_moge2.py", names) + self.assertNotIn("server/jobs.py", names) + self.assertIn("server/__init__.py", names) + self.assertIn("blendjob/environment.py", names) + self.assertIn("blendjob/installer/environment.py", names) + self.assertNotIn("install_env.py", names) self.assertNotIn("snapform/__init__.py", names) self.assertNotIn("importer.py", names) self.assertFalse( diff --git a/tests/test_service.py b/tests/test_server.py similarity index 50% rename from tests/test_service.py rename to tests/test_server.py index c0fe2e4..2946156 100644 --- a/tests/test_service.py +++ b/tests/test_server.py @@ -1,6 +1,8 @@ -import importlib.util +import importlib import sys import tempfile +import threading +import time import unittest from pathlib import Path from types import ModuleType, SimpleNamespace @@ -10,14 +12,13 @@ from unittest.mock import patch PROJECT_ROOT = Path(__file__).resolve().parents[1] -def load_service(): - specification = importlib.util.spec_from_file_location( - "snapform_service", - PROJECT_ROOT / "src" / "snapform" / "ai_service" / "server.py", - ) - module = importlib.util.module_from_spec(specification) - specification.loader.exec_module(module) - return module +def load_server_module(): + addon_directory = PROJECT_ROOT / "src" / "snapform" + sys.path.insert(0, str(addon_directory)) + try: + return importlib.import_module("server") + finally: + sys.path.remove(str(addon_directory)) class FakeCuda: @@ -33,19 +34,19 @@ class FakeCuda: return None -class ServiceTest(unittest.TestCase): +class ServerTest(unittest.TestCase): @classmethod def setUpClass(cls): - cls.service = load_service() + cls.server_module = load_server_module() def setUp(self): self.jobs_temp = tempfile.TemporaryDirectory() self.addCleanup(self.jobs_temp.cleanup) - self.jobs_directory = Path(self.jobs_temp.name) / "jobs" + self.storage_root = Path(self.jobs_temp.name) def create_cache(self): - cache = self.service.ModelSessionCache.__new__( - self.service.ModelSessionCache + cache = self.server_module.ModelSessionCache.__new__( + self.server_module.ModelSessionCache ) cache.ben2_model = None cache.ben2_key = None @@ -55,16 +56,34 @@ class ServiceTest(unittest.TestCase): cache.moge_key = None return cache + def create_server(self): + source = self.server_module.server + server = self.server_module.JobServer( + "SnapForm Job Server", storage_root=self.storage_root + ) + server.handlers.update(source.handlers) + server.resource_factories.update(source.resource_factories) + server.bind(self.storage_root) + return server + + def wait_for_job(self, context, timeout=2.0): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if context.snapshot()["state"] in {"succeeded", "failed", "cancelled"}: + return context.snapshot() + time.sleep(0.01) + self.fail(f"Job did not finish: {context.snapshot()}") + def test_ben2_model_is_reused_for_matching_runtime_key(self): loads = [] - ben2 = ModuleType("ben2_onnx") + ben2 = ModuleType("server.onnx_ben2") ben2.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) cache = self.create_cache() device = SimpleNamespace(type="cuda") - with patch.dict(sys.modules, {"ben2_onnx": ben2}): + with patch.dict(sys.modules, {"server.onnx_ben2": ben2}): first, _load_ms = cache.get_ben2("model", device) second, cached_load_ms = cache.get_ben2("model", device) @@ -76,7 +95,7 @@ class ServiceTest(unittest.TestCase): def test_da3_model_is_reused_for_matching_runtime_key(self): loads = [] - da3 = ModuleType("da3_onnx") + da3 = ModuleType("server.onnx_da3") da3.create_session = lambda model_dir, device: ( loads.append((model_dir, device)) or object() ) @@ -85,7 +104,7 @@ class ServiceTest(unittest.TestCase): with patch.dict( sys.modules, { - "da3_onnx": da3, + "server.onnx_da3": da3, }, ): first = cache.get_da3("model", device) @@ -117,34 +136,44 @@ class ServiceTest(unittest.TestCase): moge = object() - fake_moge = ModuleType("moge_onnx") + fake_moge = ModuleType("server.onnx_moge2") fake_moge.create_session = lambda _model_dir, _device: moge - with patch.dict(sys.modules, {"moge_onnx": fake_moge}): - self.assertIs(cache.get_moge_onnx("moge", "cuda"), moge) + with patch.dict(sys.modules, {"server.onnx_moge2": fake_moge}): + self.assertIs(cache.get_moge2("moge", "cuda"), moge) self.assertIsNone(cache.da3_model) self.assertIsNotNone(cache.ben2_model) - def test_dedicated_service_runs_only_one_job_at_a_time(self): - fake_models = SimpleNamespace( - snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, - clear=lambda: None, - close=lambda: None, - ) - with patch.object(self.service, "ModelManager", return_value=fake_models): - server = self.service.create_server(jobs_directory=self.jobs_directory) + def test_dedicated_server_runs_only_one_job_at_a_time(self): + server = self.server_module.JobServer("Test", storage_root=self.storage_root) + first_started = threading.Event() + release_first = threading.Event() + order = [] - parameters = {"model": "BASE"} - first = server.submit("debug-da3", parameters, job_id="job") - second = server.submit("debug-da3", parameters, job_id="other") + @server.job("example") + def example(_context, parameters): + order.append(parameters["name"]) + if parameters["name"] == "first": + first_started.set() + release_first.wait(1.0) + + first = server.submit("example", {"name": "first"}, job_id="job") + self.assertTrue(first_started.wait(1.0)) + second = server.submit("example", {"name": "second"}, job_id="other") self.assertIsNotNone(first) - self.assertIsNone(second) + self.assertIsNotNone(second) self.assertTrue(server.snapshot("instance")["busy"]) + self.assertEqual(server.snapshot("instance")["queued_jobs"], 1) self.assertNotIn("clients", server.snapshot("instance")) + release_first.set() + self.wait_for_job(first) + self.wait_for_job(second) + self.assertEqual(order, ["first", "second"]) + server.close() - def test_service_registers_tasks_instead_of_model_names(self): - server = self.service.create_server(jobs_directory=self.jobs_directory) + def test_server_registers_tasks_instead_of_model_names(self): + server = self.create_server() self.assertEqual( set(server.handlers), { @@ -161,7 +190,33 @@ class ServiceTest(unittest.TestCase): }, ) - def test_cancelling_job_keeps_service_and_model_cache_alive(self): + def test_queued_job_can_be_cancelled_before_it_runs(self): + server = self.server_module.JobServer("Test", storage_root=self.storage_root) + first_started = threading.Event() + release_first = threading.Event() + executed = [] + + @server.job("example") + def example(_context, parameters): + executed.append(parameters["name"]) + if parameters["name"] == "first": + first_started.set() + release_first.wait(1.0) + + first = server.submit("example", {"name": "first"}) + self.assertTrue(first_started.wait(1.0)) + second = server.submit("example", {"name": "second"}) + + self.assertIs(server.cancel(second.job_id), second) + release_first.set() + self.wait_for_job(first) + second_status = self.wait_for_job(second) + + self.assertEqual(second_status["state"], "cancelled") + self.assertEqual(executed, ["first"]) + server.close() + + def test_cancelling_job_keeps_server_and_model_cache_alive(self): fake_models = SimpleNamespace( snapshot=lambda: { "loaded": {"ben2": "BEN2_BASE", "da3": ""} @@ -169,36 +224,55 @@ class ServiceTest(unittest.TestCase): clear=lambda: None, close=lambda: None, ) - with patch.object(self.service, "ModelManager", return_value=fake_models): - server = self.service.create_server(jobs_directory=self.jobs_directory) + with patch.object( + self.server_module, "ModelManager", return_value=fake_models + ): + server = self.create_server() context = server.submit("remove-background", {}, job_id="job") - self.assertTrue(context.request_cancel()) + self.assertIsNotNone(server.cancel("job")) self.assertTrue(context.is_cancelled()) - server.run(context) - status = context.snapshot() + status = self.wait_for_job(context) self.assertEqual(status["state"], "cancelled") self.assertFalse(server.shutdown_event.is_set()) self.assertEqual( - server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"], + server.snapshot("instance")["resources"]["model_manager"]["loaded"]["ben2"], "BEN2_BASE", ) - def test_job_progress_is_kept_in_service_memory(self): - server = self.service.JobServer("Test", jobs_directory=self.jobs_directory) + def test_job_progress_is_kept_in_server_memory(self): + server = self.server_module.JobServer( + "Test", storage_root=self.storage_root + ) @server.job("example") def example(context, _parameters): - context.update(0.5, "Half way") + context.progress(0.5, "Half way") context = server.submit("example", {}) - server.run(context) + self.wait_for_job(context) self.assertRegex(context.job_id, r"^[0-9a-f]{32}$") self.assertEqual(context.directory.name, context.job_id) self.assertEqual(context.snapshot()["state"], "succeeded") self.assertEqual(context.snapshot()["progress"], 1.0) + def test_runtime_storage_root_owns_the_jobs_directory(self): + storage_root = self.storage_root + server = self.server_module.JobServer("Test") + server.bind(storage_root) + + @server.job("example") + def example(_context, _parameters): + return {"value": 42} + + context = server.submit("example", {}) + self.wait_for_job(context) + + self.assertEqual(server.jobs_directory, storage_root / "jobs") + self.assertEqual(context.storage_root, storage_root.resolve()) + self.assertEqual(context.snapshot()["result"], {"value": 42}) + def test_job_server_owns_and_clears_shared_resources(self): events = [] resource = SimpleNamespace( @@ -206,33 +280,40 @@ class ServiceTest(unittest.TestCase): clear=lambda: events.append("clear"), close=lambda: events.append("close"), ) - server = self.service.JobServer("Test", jobs_directory=self.jobs_directory) - server.add_resource("models", resource) + server = self.server_module.JobServer( + "Test", storage_root=self.storage_root + ) + server.add_resource("model_manager", resource) @server.job("example") def example(context, _parameters): - self.assertIs(context.resource("models"), resource) + self.assertIs(context.resource("model_manager"), resource) + release.wait(1.0) + release = threading.Event() context = server.submit("example", {}, job_id="job") - self.assertFalse(server.clear_resource("models")) - server.run(context) - self.assertTrue(server.clear_resource("models")) + self.assertFalse(server.clear_resource("model_manager")) + release.set() + self.wait_for_job(context) + self.assertTrue(server.clear_resource("model_manager")) server.close() self.assertEqual( - server.snapshot("instance")["resources"]["models"]["loaded"], + server.snapshot("instance")["resources"]["model_manager"]["loaded"], "model", ) self.assertEqual(events, ["clear", "close"]) - def test_service_dispatches_frame_extraction_without_models(self): + def test_server_dispatches_frame_extraction_without_models(self): fake_models = SimpleNamespace( snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, clear=lambda: None, close=lambda: None, ) - with patch.object(self.service, "ModelManager", return_value=fake_models): - server = self.service.create_server(jobs_directory=self.jobs_directory) + with patch.object( + self.server_module, "ModelManager", return_value=fake_models + ): + server = self.create_server() with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_directory = root / "frames" @@ -248,8 +329,7 @@ class ServiceTest(unittest.TestCase): }, job_id="job", ) - server.run(context) - status = context.snapshot() + status = self.wait_for_job(context) self.assertTrue( (context.directory / "frame.0001.jpg").is_file() ) diff --git a/tests/test_worker.py b/tests/test_worker.py index 5f94329..99edf56 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -1,7 +1,5 @@ import importlib.util -import json import os -import shutil import sys import tempfile import unittest @@ -16,14 +14,13 @@ import numpy as np PROJECT_ROOT = Path(__file__).resolve().parents[1] -def load_worker(): - specification = importlib.util.spec_from_file_location( - "snapform_worker", - PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py", - ) - module = importlib.util.module_from_spec(specification) - specification.loader.exec_module(module) - return module +def load_server_module(name): + addon_directory = PROJECT_ROOT / "src" / "snapform" + sys.path.insert(0, str(addon_directory)) + try: + return __import__(f"server.{name}", fromlist=[name]) + finally: + sys.path.remove(str(addon_directory)) def _make_fake_av(frame_count): @@ -69,10 +66,54 @@ def _make_fake_av(frame_count): return FakeAv() +class FakeModels: + def __init__(self, directory="model"): + self.model_directory = Path(directory) + + def directory(self, _model): + return self.model_directory + + def get_ben2(self, _model_directory, _device): + from server.onnx_ben2 import create_session + + return create_session(self.model_directory, "cpu"), 0.0 + + def get_da3(self, _model_directory, _device): + return object() + + +class FakeJobContext: + def __init__(self, directory, job_type, models=None): + self.directory = Path(directory) + self.job_type = job_type + self.models = models or FakeModels() + self.status = {} + + def resource(self, name): + if name != "model_manager": + raise KeyError(name) + return self.models + + def progress(self, progress, message, **details): + self.status = { + "progress": progress, + "message": message, + **details, + } + + def check_cancelled(self): + return None + + class WorkerTest(unittest.TestCase): @classmethod def setUpClass(cls): - cls.worker = load_worker() + cls.media = load_server_module("media") + cls.geometry = load_server_module("geometry") + cls.background = load_server_module("background") + cls.moge2 = load_server_module("moge2") + cls.da3 = load_server_module("da3") + cls.download = load_server_module("model_download") def test_foreground_preserves_source_rgb_and_ben2_alpha(self): alpha = SimpleNamespace(size=(2, 1)) @@ -88,7 +129,7 @@ class WorkerTest(unittest.TestCase): ) foreground = SimpleNamespace(convert=lambda _mode: foreground_rgba) - result = self.worker.preserve_source_rgb(source, foreground) + result = self.background.preserve_source_rgb(source, foreground) self.assertIs(result, source_rgba) self.assertIs(result.applied_alpha, alpha) @@ -100,62 +141,22 @@ class WorkerTest(unittest.TestCase): (root / "a.jpg").touch() (root / "notes.txt").touch() - images = self.worker.collect_images(root) + images = self.media.collect_images(root) self.assertEqual([path.name for path in images], ["a.jpg", "b.PNG"]) def test_ai_tasks_support_cooperative_cancellation(self): - with self.assertRaises(self.worker.TaskCancelled): - self.worker.check_cancelled(lambda: True) + with self.assertRaises(self.media.TaskCancelled): + self.media.check_cancelled(lambda: True) - self.worker.check_cancelled(lambda: False) - - def test_reconstruct_parser_accepts_background_model_directory(self): - args = self.worker.parser().parse_args( - [ - "reconstruct", - "--input", - "input.png", - "--model-dir", - "da3", - "--output-dir", - "result", - "--export-format", - "mini_npz", - "--status", - "status.json", - "--remove-background-model-dir", - "ben2", - "--process-res", - "1024", - "--process-res-method", - "upper_bound_resize", - "--conf-thresh-percentile", - "40", - "--num-max-points", - "1000", - "--show-cameras", - "0", - "--gs-views-interval", - "1", - "--video-frames", - "16", - "--device", - "auto", - "--use-ray-pose", - "1", - "--ref-view-strategy", - "first", - ] - ) - self.assertEqual(args.remove_background_model_dir, "ben2") + self.media.check_cancelled(lambda: False) def test_prepares_a_single_image_file(self): with tempfile.TemporaryDirectory() as directory: image_path = Path(directory) / "input.png" image_path.touch() - images, temporary_directory = self.worker.prepare_input( + images, temporary_directory = self.media.prepare_input( image_path, 16, directory, @@ -171,7 +172,7 @@ class WorkerTest(unittest.TestCase): (root / "frame-0002.png").touch() (root / "notes.txt").touch() - images, temporary_directory = self.worker.prepare_input( + images, temporary_directory = self.media.prepare_input( root, 16, directory, @@ -189,7 +190,7 @@ class WorkerTest(unittest.TestCase): for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() - images, temporary_directory = self.worker.prepare_input( + images, temporary_directory = self.media.prepare_input( root, 3, directory, @@ -207,7 +208,7 @@ class WorkerTest(unittest.TestCase): for index in range(1, 6): (root / f"frame-{index:04d}.jpg").touch() - images, temporary_directory = self.worker.prepare_input( + images, temporary_directory = self.media.prepare_input( root, 3, directory, @@ -224,7 +225,7 @@ class WorkerTest(unittest.TestCase): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) - image_paths = self.worker.extract_video_frames( + image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, @@ -242,7 +243,7 @@ class WorkerTest(unittest.TestCase): with tempfile.TemporaryDirectory() as directory: root = Path(directory) fake_av = _make_fake_av(100) - image_paths = self.worker.extract_video_frames( + image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, @@ -260,7 +261,7 @@ class WorkerTest(unittest.TestCase): def test_extract_video_frames_tolerates_short_video_tail(self): with tempfile.TemporaryDirectory() as directory: root = Path(directory) - image_paths = self.worker.extract_video_frames( + image_paths = self.media.extract_video_frames( root / "clip.mkv", root / "out", 10, @@ -280,11 +281,11 @@ class WorkerTest(unittest.TestCase): valid_path, depth=np.zeros((4, 4)), ) - self.assertTrue(self.worker.npz_is_complete(valid_path)) + self.assertTrue(self.media.npz_is_complete(valid_path)) empty_path = root / "empty.npz" empty_path.write_bytes(b"PK\x05\x06" + b"\x00" * 18) - self.assertFalse(self.worker.npz_is_complete(empty_path)) + self.assertFalse(self.media.npz_is_complete(empty_path)) def test_wait_for_exported_paths_waits_for_complete_npz(self): with tempfile.TemporaryDirectory() as directory: @@ -293,11 +294,11 @@ class WorkerTest(unittest.TestCase): npz_path.parent.mkdir(parents=True) npz_path.touch() with patch.object( - self.worker, + self.media, "npz_is_complete", side_effect=[False, True], ): - result = self.worker.wait_for_exported_paths( + result = self.media.wait_for_exported_paths( root, "mini_npz", timeout=2.0, @@ -324,10 +325,10 @@ class WorkerTest(unittest.TestCase): return original_replace(path, target) with patch.object(path_type, "replace", new=locked_once), patch.object( - self.worker.time, + self.media.time, "sleep", ) as sleep: - result = self.worker.move_exported_file( + result = self.media.move_exported_file( source, destination, ) @@ -342,7 +343,7 @@ class WorkerTest(unittest.TestCase): input_path.touch() with self.assertRaisesRegex(RuntimeError, "Unsupported input type"): - self.worker.prepare_input(input_path, 16, directory) + self.media.prepare_input(input_path, 16, directory) def test_download_model_uses_actual_request_without_preflight(self): calls = [] @@ -353,7 +354,7 @@ class WorkerTest(unittest.TestCase): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", - status="status.json", + progress=lambda *_args, **_kwargs: None, allow_pattern=[], ) @@ -361,17 +362,16 @@ class WorkerTest(unittest.TestCase): sys.modules, {"huggingface_hub": fake_hugging_face}, ), patch.object( - self.worker, + self.download, "create_download_progress_class", return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object( - self.worker, + self.download, "_verify_model_complete", return_value=None, ): args.model_dir = str(Path(directory) / "models" / "DA3") - args.status = str(Path(directory) / "status.json") - self.worker.download_model(args) + self.download.download_model(args) self.assertEqual(len(calls), 1) @@ -381,7 +381,7 @@ class WorkerTest(unittest.TestCase): args = SimpleNamespace( repository="PramaLLC/BEN2", model_dir="model", - status="status.json", + progress=lambda *_args, **_kwargs: None, allow_pattern=[], ) @@ -389,47 +389,45 @@ class WorkerTest(unittest.TestCase): sys.modules, {"huggingface_hub": fake_hugging_face}, ), patch.object( - self.worker, + self.download, "create_download_progress_class", return_value=object, ), patch.dict(os.environ, {"HF_ENDPOINT": ""}): 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) + self.download.download_model(args) def test_download_model_reports_all_failed_sources(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", - status="status.json", + progress=lambda *_args, **_kwargs: None, allow_pattern=[], r2_directory="DA3Small", r2_file=["config.json|100|abc"], ) with tempfile.TemporaryDirectory() as directory, patch.object( - self.worker, + self.download, "download_from_hf", side_effect=OSError("direct down"), ), patch.object( - self.worker, + self.download, "download_from_r2", side_effect=OSError("mirror down"), ): args.model_dir = str(Path(directory) / "model") - args.status = str(Path(directory) / "status.json") with self.assertRaisesRegex(RuntimeError, "models.omoolab.xyz"): - self.worker.download_model(args) + self.download.download_model(args) def test_download_model_prefers_direct_hf_then_r2(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", - status="status.json", + progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", r2_file=["config.json|100|abc"], @@ -444,83 +442,28 @@ class WorkerTest(unittest.TestCase): order.append("r2") with tempfile.TemporaryDirectory() as directory, patch.object( - self.worker, + self.download, "download_from_hf", side_effect=fake_download_from_hf, ), patch.object( - self.worker, + self.download, "download_from_r2", side_effect=fake_download_from_r2, ), patch.object( - self.worker, + self.download, "_verify_model_complete", return_value=None, ): args.model_dir = str(Path(directory) / "models" / "DA3") - args.status = str(Path(directory) / "status.json") - self.worker.download_model(args) + self.download.download_model(args) self.assertEqual(order, ["hf", "r2"]) - def test_download_declared_skips_hf_when_r2_preferred(self): - root = PROJECT_ROOT / ".tmp-models-from-r2" - model_dir = root / "models" / "DA3" - args = SimpleNamespace( - repository="depth-anything/DA3-SMALL", - model_dir=str(model_dir), - status="status.json", - allow_pattern=[], - r2_directory="DA3Small", - r2_file=["config.json|100|abc"], - ) - order = [] - try: - self.worker._mark_models_from_r2(model_dir) - with patch.object( - self.worker, - "download_from_hf", - side_effect=lambda *a, **k: order.append("hf"), - ), patch.object( - self.worker, - "download_from_r2", - side_effect=lambda *a, **k: order.append("r2"), - ), patch.object( - self.worker, - "_verify_model_complete", - return_value=None, - ), patch.object(self.worker, "write_status"): - self.worker._download_declared_or_fail( - args, model_dir, "DA3 Small", 0.0, 1.0, {} - ) - finally: - shutil.rmtree(root, ignore_errors=True) - - self.assertEqual(order, ["r2"]) - - def test_models_from_r2_roundtrip_preserves_fields(self): - root = PROJECT_ROOT / ".tmp-models-from-r2" - model_dir = root / "models" / "DA3" - runtime_state = root / "runtime.json" - try: - runtime_state.parent.mkdir(parents=True, exist_ok=True) - runtime_state.write_text( - json.dumps({"keep": "value"}), - encoding="utf-8", - ) - self.assertFalse(self.worker._models_from_r2(model_dir)) - self.worker._mark_models_from_r2(model_dir) - self.assertTrue(self.worker._models_from_r2(model_dir)) - payload = json.loads(runtime_state.read_text(encoding="utf-8")) - self.assertEqual(payload["keep"], "value") - self.assertTrue(payload["models_from_r2"]) - finally: - shutil.rmtree(root, ignore_errors=True) - def test_download_from_r2_downloads_declared_files(self): args = SimpleNamespace( repository="depth-anything/DA3-SMALL", model_dir="model", - status="status.json", + progress=lambda *_args, **_kwargs: None, allow_pattern=["config.json"], r2_directory="DA3Small", r2_file=["config.json|100|abc"], @@ -528,13 +471,12 @@ class WorkerTest(unittest.TestCase): downloaded = [] with tempfile.TemporaryDirectory() as directory, patch.object( - self.worker, + self.download, "_download_file_from_url", side_effect=lambda *args, **kwargs: downloaded.append(args[1].name), - ), patch.object(self.worker, "write_status"): + ): args.model_dir = str(Path(directory) / "model") - args.status = str(Path(directory) / "status.json") - self.worker.download_from_r2(args, args.model_dir) + self.download.download_from_r2(args, args.model_dir) self.assertEqual(downloaded, ["config.json"]) @@ -589,7 +531,7 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("ben2_onnx") + fake_ben2 = ModuleType("server.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() @@ -601,22 +543,18 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "ben2_onnx": fake_ben2, + "server.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): root = Path(directory) input_path = root / "input.png" input_path.touch() - args = SimpleNamespace( - input=str(input_path), - model_dir="model", - output_dir=str(root / "result"), - status=str(root / "status.json"), - device="cpu", + context = FakeJobContext(root / "result", "remove-background") + self.background.remove_background( + context, + {"input": str(input_path), "model": "BEN2", "device": "cpu"}, ) - self.worker.remove_background(args) - status = json.loads(Path(args.status).read_text(encoding="utf-8")) self.assertTrue((root / "result" / "foreground.png").is_file()) @@ -628,7 +566,7 @@ class WorkerTest(unittest.TestCase): save_options[0], {"format": "PNG", "compress_level": 1}, ) - self.assertIn("Background removed in", status["message"]) + self.assertIn("Background removed in", context.status["message"]) def test_ben2_background_removal_processes_each_frame(self): inference_options = [] @@ -677,7 +615,7 @@ class WorkerTest(unittest.TestCase): ) fake_torch.device = lambda device_type: SimpleNamespace(type=device_type) fake_torch.inference_mode = nullcontext - fake_ben2 = ModuleType("ben2_onnx") + fake_ben2 = ModuleType("server.onnx_ben2") fake_ben2.create_session = lambda _model_dir, _device: object() fake_ben2.infer = lambda _model, image: ( inference_options.append((image, {})) or FakeImage() @@ -689,7 +627,7 @@ class WorkerTest(unittest.TestCase): sys.modules, { "torch": fake_torch, - "ben2_onnx": fake_ben2, + "server.onnx_ben2": fake_ben2, "PIL": fake_pil, }, ): @@ -698,17 +636,18 @@ class WorkerTest(unittest.TestCase): 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, + context = FakeJobContext(root / "result", "remove-background") + self.background.remove_background( + context, + { + "input": str(input_directory), + "model": "BEN2", + "device": "cpu", + "video_frames": 8, + }, ) - self.worker.remove_background(args) - result_directory = Path(args.output_dir) + result_directory = context.directory self.assertTrue( (result_directory / "foreground.0001.png").is_file() ) @@ -726,49 +665,6 @@ class WorkerTest(unittest.TestCase): for _image, options in inference_options: self.assertEqual(options, {}) - 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) @@ -776,17 +672,17 @@ class WorkerTest(unittest.TestCase): 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, + context = FakeJobContext(root / "result", "extract-frames") + self.media.extract_frames( + context, + { + "input": str(input_directory), + "video_frames": 3, + "start_frame": 2, + }, ) - self.worker.extract_frames(args) - - result_directory = Path(args.output_dir) + result_directory = context.directory self.assertEqual( sorted( path.name @@ -795,11 +691,8 @@ class WorkerTest(unittest.TestCase): ), ["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"]) + self.assertEqual(context.status["progress"], 1.0) + self.assertIn("3 frame", context.status["message"]) def test_uses_da3_export_parameters_for_formats_and_debug_mesh(self): inference_options = [] @@ -812,7 +705,7 @@ class WorkerTest(unittest.TestCase): intrinsics=np.eye(3, dtype=np.float32)[None], processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8), ) - fake_da3 = ModuleType("da3_onnx") + fake_da3 = ModuleType("server.onnx_da3") fake_da3.create_session = lambda _model_dir, _device: object() def fake_infer(_session, images, **options): @@ -820,7 +713,7 @@ class WorkerTest(unittest.TestCase): return prediction fake_da3.infer = fake_infer - fake_runtime = ModuleType("onnx_runtime") + fake_runtime = ModuleType("server.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" def fake_mesh_export( @@ -865,23 +758,23 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "da3_onnx": fake_da3, - "onnx_runtime": fake_runtime, + "server.onnx_da3": fake_da3, + "server.onnx_runtime": fake_runtime, }, ), patch.object( - self.worker, + self.da3, "export_textured_mesh_glb", side_effect=fake_mesh_export, ), patch.object( - self.worker, + self.da3, "export_da3_point_cloud", side_effect=fake_point_cloud, ), patch.object( - self.worker, + self.da3, "export_da3_depth_preview", side_effect=fake_depth_preview, ), patch.object( - self.worker, + self.da3, "export_da3_mini_npz", side_effect=fake_mini_npz, ): @@ -892,32 +785,31 @@ class WorkerTest(unittest.TestCase): "mesh", ) for export_format in debug_formats: - args = SimpleNamespace( - input=str(input_path), - video_frames=16, - output_dir=str(root / export_format), - export_format=export_format, - status=str(root / f"{export_format}.json"), - device="cpu", - model_dir="model", - process_res=518, - process_res_method="lower_bound_resize", - use_ray_pose=1, - ref_view_strategy="middle", - conf_thresh_percentile=35.0, - num_max_points=12345, - show_cameras=0, - gs_views_interval=2, - mesh_filter_edges=True, - mesh_max_edge_length=0.2, + context = FakeJobContext( + root / export_format, + "debug-da3", ) - self.worker.reconstruct(args) - status = json.loads( - Path(args.status).read_text(encoding="utf-8") + self.da3.generate_da3( + context, + { + "input": str(input_path), + "model": "DA3", + "video_frames": 16, + "export_format": export_format, + "device": "cpu", + "process_res": 518, + "process_res_method": "lower_bound_resize", + "use_ray_pose": 1, + "ref_view_strategy": "middle", + "conf_thresh_percentile": 35.0, + "num_max_points": 12345, + "mesh_filter_edges": True, + "mesh_max_edge_length": 0.2, + }, ) - self.assertEqual(status["progress"], 1.0) - self.assertIn(export_format, status["message"]) - result_directory = Path(args.output_dir) + self.assertEqual(context.status["progress"], 1.0) + self.assertIn(export_format, context.status["message"]) + result_directory = context.directory if export_format == "depth_vis": self.assertTrue((result_directory / "0000.jpg").is_file()) self.assertFalse((result_directory / "depth_vis").exists()) @@ -938,7 +830,7 @@ class WorkerTest(unittest.TestCase): def test_reconstruct_copies_frames_for_multi_frame_input(self): inference_inputs = [] - fake_da3 = ModuleType("da3_onnx") + fake_da3 = ModuleType("server.onnx_da3") fake_da3.create_session = lambda _model_dir, _device: object() def fake_infer(_session, images, **_options): @@ -952,7 +844,7 @@ class WorkerTest(unittest.TestCase): ) fake_da3.infer = fake_infer - fake_runtime = ModuleType("onnx_runtime") + fake_runtime = ModuleType("server.onnx_runtime") fake_runtime.provider_label = lambda _session: "CPU" with tempfile.TemporaryDirectory() as directory: @@ -964,32 +856,33 @@ class WorkerTest(unittest.TestCase): with patch.dict( sys.modules, { - "da3_onnx": fake_da3, - "onnx_runtime": fake_runtime, + "server.onnx_da3": fake_da3, + "server.onnx_runtime": fake_runtime, }, ): - 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="", - inference_mode="per_frame", + context = FakeJobContext( + root / "result", + "generate-da3-depth", + ) + self.da3.generate_da3( + context, + { + "input": str(input_directory), + "model": "DA3", + "video_frames": 16, + "export_format": "mini_npz", + "device": "cpu", + "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, + "inference_mode": "per_frame", + }, ) - self.worker.reconstruct(args) - result_directory = Path(args.output_dir) + result_directory = context.directory self.assertTrue( (result_directory / "frame.0001.jpg").is_file() ) @@ -1006,14 +899,13 @@ class WorkerTest(unittest.TestCase): self.assertFalse( (result_directory / "exports").exists() ) - status = json.loads(Path(args.status).read_text(encoding="utf-8")) - self.assertEqual(status["progress"], 1.0) + self.assertEqual(context.status["progress"], 1.0) def test_mesh_confidence_uses_percentile(self): depth = np.ones((1, 2, 2), dtype=np.float32) confidence = np.array([[[1.0, 2.0], [3.0, 4.0]]]) - threshold = self.worker.confidence_percentile_threshold( + threshold = self.geometry.confidence_percentile_threshold( depth, confidence, 50.0, @@ -1027,12 +919,12 @@ class WorkerTest(unittest.TestCase): x, y = np.meshgrid([0.0, 1.0], [0.0, 1.0]) points = np.stack((x, y, depth), axis=-1) - threshold = self.worker.confidence_percentile_threshold( + threshold = self.geometry.confidence_percentile_threshold( depth, confidence, 0.0, ) - faces = self.worker.mesh_triangle_faces( + faces = self.geometry.mesh_triangle_faces( points, depth, confidence, @@ -1054,7 +946,7 @@ class WorkerTest(unittest.TestCase): ) points = np.stack((x, y, depth), axis=-1) - faces = self.worker.mesh_triangle_faces( + faces = self.geometry.mesh_triangle_faces( points, depth, confidence, @@ -1076,7 +968,7 @@ class WorkerTest(unittest.TestCase): foreground = np.ones((2, 3), dtype=bool) foreground[:, 0] = False - faces = self.worker.mesh_triangle_faces( + faces = self.geometry.mesh_triangle_faces( points, depth, confidence, @@ -1096,12 +988,12 @@ class WorkerTest(unittest.TestCase): discontinuous = smooth.copy() discontinuous[:, 3:] += 1.0 - self.assertFalse(self.worker.depth_edge_mask(smooth).any()) - self.assertTrue(self.worker.depth_edge_mask(discontinuous)[:, 2:4].all()) + self.assertFalse(self.geometry.depth_edge_mask(smooth).any()) + self.assertTrue(self.geometry.depth_edge_mask(discontinuous)[:, 2:4].all()) def test_higher_edge_culling_allows_less_stretch(self): limits = [ - self.worker.edge_culling_stretch_limit(value) + self.geometry.edge_culling_stretch_limit(value) for value in (0.0, 50.0, 100.0) ] @@ -1119,17 +1011,17 @@ class WorkerTest(unittest.TestCase): stretched_points[0, 1] = (2.0, 0.0, 2.0) triangles = np.asarray([[0, 2, 3], [0, 3, 1]], dtype=np.int32) - flat_mask = self.worker.triangle_edge_stretch_mask( + flat_mask = self.geometry.triangle_edge_stretch_mask( flat_points, triangles, 1.0, ) - loose_mask = self.worker.triangle_edge_stretch_mask( + loose_mask = self.geometry.triangle_edge_stretch_mask( stretched_points, triangles, 2.0, ) - strict_mask = self.worker.triangle_edge_stretch_mask( + strict_mask = self.geometry.triangle_edge_stretch_mask( stretched_points, triangles, 1.0, @@ -1151,7 +1043,7 @@ class WorkerTest(unittest.TestCase): confidence = np.ones((2, 3), dtype=np.float32) confidence[:, 0] = 0.0 - faces = self.worker.mesh_triangle_faces( + faces = self.geometry.mesh_triangle_faces( points, depth, confidence, @@ -1178,12 +1070,12 @@ class WorkerTest(unittest.TestCase): points = np.stack((x, y, depth), axis=-1) face_counts = [] for percentile in (0.0, 50.0, 100.0): - threshold = self.worker.confidence_percentile_threshold( + threshold = self.geometry.confidence_percentile_threshold( depth, confidence, percentile, ) - faces = self.worker.mesh_triangle_faces( + faces = self.geometry.mesh_triangle_faces( points, depth, confidence, @@ -1199,7 +1091,7 @@ class WorkerTest(unittest.TestCase): vertices = np.arange(18, dtype=np.float32).reshape(6, 3) faces = np.asarray([[1, 3, 5]], dtype=np.int32) - compact_vertices, compact_faces = self.worker.compact_mesh( + compact_vertices, compact_faces = self.geometry.compact_mesh( vertices, faces, ) @@ -1220,7 +1112,7 @@ class WorkerTest(unittest.TestCase): dtype=bool, ) - filtered = self.worker.remove_background_faces(faces, foreground) + filtered = self.geometry.remove_background_faces(faces, foreground) np.testing.assert_array_equal(filtered, faces[2:]) @@ -1258,14 +1150,14 @@ class WorkerTest(unittest.TestCase): sys.modules, {"PIL": fake_pil}, ), patch.object( - self.worker, + self.geometry, "compact_mesh", side_effect=lambda vertices, faces: ( vertices, faces, ), ): - result_path, texture_path = self.worker.export_image_mesh_data( + result_path, texture_path = self.geometry.export_image_mesh_data( prediction, Path(directory) / "input.png", Path(directory) / "result", @@ -1313,7 +1205,7 @@ class WorkerTest(unittest.TestCase): ) with tempfile.TemporaryDirectory() as directory: - output_path = self.worker.export_prediction_attributes( + output_path = self.geometry.export_prediction_attributes( prediction, directory, )