From b3a07cde320327de0108efefc486352fe0dd599e Mon Sep 17 00:00:00 2001 From: nan Date: Sun, 16 Aug 2026 21:04:10 +0800 Subject: [PATCH] Add reusable local HTTP job service for Blender add-ons - Implement JobClient for handling HTTP requests to the JobServer. - Create ModalJobOperatorBase for managing Blender job lifecycle. - Introduce JobContext and JobServer for job management and state handling. - Add ServiceConnection and ServiceController for managing service processes. - Implement job submission, cancellation, and resource management endpoints in the server. - Ensure proper error handling and job state updates throughout the service. --- README.md | 8 +- docs/blendjob.md | 179 +++ docs/development.md | 19 +- docs/features.md | 6 +- docs/index.md | 1 + docs/service.md | 44 +- docs/surface.md | 45 +- src/snapform/__init__.py | 10 +- src/snapform/ai_service/__init__.py | 1 + src/snapform/{ => ai_service}/ben2_onnx.py | 0 src/snapform/{ => ai_service}/da3_onnx.py | 0 .../{worker.py => ai_service/jobs.py} | 40 +- src/snapform/ai_service/model_manager.py | 153 +++ src/snapform/{ => ai_service}/moge_onnx.py | 0 src/snapform/{ => ai_service}/onnx_runtime.py | 0 src/snapform/ai_service/server.py | 145 +++ src/snapform/blender/__init__.py | 1 + .../depth_plane.py} | 55 +- .../{image_mesh.py => blender/surface.py} | 57 +- src/snapform/blendjob/__init__.py | 37 + src/snapform/blendjob/client.py | 99 ++ src/snapform/blendjob/operator.py | 267 ++++ src/snapform/blendjob/server.py | 350 +++++ src/snapform/blendjob/service.py | 208 +++ src/snapform/client.py | 841 +++--------- src/snapform/install_env.py | 2 +- src/snapform/operators.py | 1141 +++++++---------- src/snapform/panel.py | 32 +- src/snapform/preferences.py | 5 +- src/snapform/service.py | 598 --------- tests/test_addon.py | 657 +++------- tests/test_ben2_onnx.py | 7 +- tests/test_da3_onnx.py | 7 +- tests/test_download_progress.py | 2 +- tests/test_image_depth.py | 6 +- tests/test_moge_onnx.py | 2 +- tests/test_pack.py | 7 +- tests/test_service.py | 211 ++- tests/test_worker.py | 14 +- tools/validate_depth_plane_asset.py | 4 +- tools/validate_geometry_nodes_asset.py | 4 +- tools/validate_moge_depth_plane.py | 2 +- tools/validate_moge_surface.py | 4 +- 43 files changed, 2600 insertions(+), 2671 deletions(-) create mode 100644 docs/blendjob.md create mode 100644 src/snapform/ai_service/__init__.py rename src/snapform/{ => ai_service}/ben2_onnx.py (100%) rename src/snapform/{ => ai_service}/da3_onnx.py (100%) rename src/snapform/{worker.py => ai_service/jobs.py} (98%) create mode 100644 src/snapform/ai_service/model_manager.py rename src/snapform/{ => ai_service}/moge_onnx.py (100%) rename src/snapform/{ => ai_service}/onnx_runtime.py (100%) create mode 100644 src/snapform/ai_service/server.py create mode 100644 src/snapform/blender/__init__.py rename src/snapform/{image_depth.py => blender/depth_plane.py} (98%) rename src/snapform/{image_mesh.py => blender/surface.py} (94%) create mode 100644 src/snapform/blendjob/__init__.py create mode 100644 src/snapform/blendjob/client.py create mode 100644 src/snapform/blendjob/operator.py create mode 100644 src/snapform/blendjob/server.py create mode 100644 src/snapform/blendjob/service.py delete mode 100644 src/snapform/service.py diff --git a/README.md b/README.md index 297d75b..830227b 100644 --- a/README.md +++ b/README.md @@ -75,10 +75,7 @@ Session 逐帧推理，显存峰值不随视频帧数平方增长；Debug 多帧 ## 任务结果保存在哪里 -任务结果保存在 `~/.snapform/jobs/` 下，每个任务一个文件夹，名称统一为 -`动作-对象_日期-时间-微秒`。Packed Image 和序列帧的输入副本会直接归档到对应任务根目录， -不再使用公共 `inputs/` 文件夹；导出产物也尽量平铺。视频或序列帧生成的面片会引用任务文件夹里的 -序列帧，请保留对应文件夹，不要随意删除。 +任务结果保存在 `~/.snapform/jobs/` 下。Job Server 接受任务时自动生成 opaque Job ID，并创建同名文件夹；Blender 不预先申请 ID 或目录。Packed Image 和序列帧的临时输入由 Blender 在任务结束后清理，Service 产物直接写入该 Job 目录。视频或序列帧生成的面片会引用其中的序列帧，请保留对应文件夹，不要随意删除。 ## 常见问题 @@ -90,4 +87,7 @@ Session 逐帧推理，显存峰值不随视频帧数平方增长；Debug 多帧 ## 技术文档 +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 目录。 + 详细功能与技术内容见 [docs/index.md](docs/index.md)。 +可复用的本地 Job Service API 与插件接入方式见 [docs/blendjob.md](docs/blendjob.md)。 diff --git a/docs/blendjob.md b/docs/blendjob.md new file mode 100644 index 0000000..636b442 --- /dev/null +++ b/docs/blendjob.md @@ -0,0 +1,179 @@ +# BlendJob 开发者指南 + +## 命名与边界 + +通用 module 命名为 `blendjob`。这个名字表达的是 Blender 插件与独立 Job 进程之间的桥接，不绑定 AI、模型框架或 SnapForm。 + +- `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` + +`blendjob` 不知道模型、输出文件或 Blender Operator。插件专属代码负责注册 Job 和 Resource，并把参数交给自己的执行逻辑。 + +## 最小接入 + +服务端只创建 Server 并注册业务函数： + +```python +server = JobServer("Example", jobs_directory="D:/example/jobs") + + +@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)) +``` + +Blender 端直接继承一个基类。普通任务只需要声明 Controller、Job Type 和怎么传值： + +```python +class GenerateDepth(ModalJobOperatorBase): + bl_idname = "example.generate_depth" + bl_label = "Generate Depth" + controller_factory = staticmethod(service_controller) + job_type = "generate-depth" + + def job_parameters(self, context): + return {"input": context.scene.source_path} +``` + +这个 Operator 已经具备异步提交、Server 生成 Job ID 与目录、HTTP 轮询、进度、Esc 取消、错误报告和 Timer 清理。只有需要把产物导回 Blender 时，才覆盖 `on_job_succeeded(context, job, status)`；产物目录是 `job.directory`。 + +## Resource + +Job 表示一次有终态的操作；Resource 表示在 Service 生命周期内被多个 Job 共用的状态。Resource 可以按需实现 `snapshot()`、`clear()` 和 `close()`： + +```python +models = server.add_resource("models", ModelManager(models_directory)) + + +@server.job("generate-depth") +def generate_depth(context, parameters): + models = context.resource("models") + session = models.get(parameters["model"], parameters["device"]) +``` + +`JobServer` 在同一把调度锁下阻止运行中清理 Resource，关闭时按注册逆序调用 `close()`。Resource Snapshot 自动进入 `/health` 的 `resources` 字段。 + +## 服务端使用 + +开发者先创建 `JobServer`，再用装饰器注册 Job 类型： + +```python +from blendjob import JobCancelled, JobServer + + +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, + ) +``` + +处理函数签名固定为 `handler(context, parameters)`。`parameters` 必须可由 JSON 表示。客户端提交的 `JobRequest` 不包含 ID；`JobServer` 接受请求时自动生成 opaque Job ID，并创建 `//`。处理函数通过 `context.job_id` 与 `context.directory` 获取它们。客户端不创建 Job 目录，也不应把输出目录作为业务参数传入。Job API 只保存内存状态和少量结果元数据，不创建 `status.json`。 + +`context.update()` 支持以下通用字段： + +- `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"]) +``` + +专属 Service 进程由 `ServiceController` 管理。插件只提供启动命令、工作目录和日志路径，端口、Instance ID、进程句柄、连接验证以及 `start()`、`ensure()`、`enable()`、`restart()`、`stop()`、`detach()` 都由 Controller 负责： + +```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() +``` + +对应 HTTP API： + +- `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`：关闭当前专属服务 + +Service 只监听 `127.0.0.1`，专属实例由启动它的 `ServiceController` 保存随机端口与 Instance ID。这个本地专属进程协议不额外使用 Token 鉴权。 + +`ModalJobOperatorBase` 默认使用 Blender Window Manager 的进度接口。插件若有自己的状态栏，可覆盖 `update_job_ui()` 与 `redraw_job_ui()`；这些属于界面定制，不是接入 Job 所必需的代码。运行态使用通用 `ModalJobState`，只包含 Controller、Job ID、目录、进度和取消状态。插件特有的完成参数可由 `create_job_data()` 返回，并通过 `job.data` 读取，不应加入通用状态类。 + +## 生命周期建议 + +每个 Blender 实例为每个插件启动一个专属 Job Server。随机端口与 Instance ID 由 `ServiceController` 生成并只保存在内存中，不写公共发现文件，也不扫描其它本地服务。 + +启动服务时传入 Blender 的 PID。服务同时监听显式 `/shutdown` 和父进程存活状态，因此正常禁用插件时可立即停止，Blender 异常退出后也能自行退出。这个方案允许同一机器出现多个服务进程，换取明确的所有权和插件隔离。 + +模型与 Environment 可以继续放在插件自己的共享 Storage Root 中；共享磁盘缓存不代表共享服务进程。若多个 Blender 可能同时写同一个模型目录，插件仍需对下载或更新操作单独加文件锁。 + +## 在其它插件复用 + +当前实现位于 `src/snapform/blendjob/`，没有 SnapForm 依赖。复用时复制整个 `blendjob` package，并在插件专属入口中完成以下工作： + +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 扩展增加安装依赖。 diff --git a/docs/development.md b/docs/development.md index 5e620c9..ae32aa8 100644 --- a/docs/development.md +++ b/docs/development.md @@ -15,16 +15,29 @@ uv run python pack.py ├── src/ │ └── snapform/ │ ├── blender_manifest.toml -│ ├── service.py │ ├── install_env.py │ ├── models.py -│ ├── image_depth.py │ ├── operators.py │ ├── panel.py │ ├── preferences.py │ ├── properties.py │ ├── client.py -│ └── worker.py +│ ├── blendjob/ +│ │ ├── client.py +│ │ ├── operator.py +│ │ ├── service.py +│ │ └── server.py +│ ├── blender/ +│ │ ├── depth_plane.py +│ │ └── surface.py +│ └── ai_service/ +│ ├── server.py +│ ├── model_manager.py +│ ├── jobs.py +│ ├── ben2_onnx.py +│ ├── da3_onnx.py +│ ├── moge_onnx.py +│ └── onnx_runtime.py └── tests/ ``` diff --git a/docs/features.md b/docs/features.md index 13ea5ce..628113d 100644 --- a/docs/features.md +++ b/docs/features.md @@ -58,7 +58,7 @@ （Image Empty 右键菜单）使用的推理设备，默认取本机最快的 DirectML/CoreML provider，否则 CPU，且只列出 本机支持的设备。 4. 在 `AI Service` 点击 `Install Environment`。首次安装会下载 Python、ONNX Runtime、FastAPI、Uvicorn； - uv 已内置在扩展中，DA3、BEN2 与 MoGe 不安装源码包， + 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 已存在时省略其阶段并重新编号。 @@ -71,8 +71,8 @@ Environment 安装阶段与 `Models` 面板单独下载使用相同的模型名称格式，不会改写为通用的 `Downloading packages`。 Environment 安装后即使模型下载中断也不影响 Service 使用，可稍后在 `Models` 面板单独补下。 安装不做启动网络预检：依赖安装默认依次尝试清华、阿里镜像并以 pypi.org 兜底；模型下载由 - worker 自行探测回退到 OmooLab 镜像，切换结果写入 service.json 的 models_from_r2 字段， - 本次周期内后续模型直接走 R2；需要手动指定依赖镜像时可通过 + Downloader 自行探测回退到 OmooLab 镜像，切换结果写入 `runtime.json` 的 `models_from_r2` 字段， + 后续下载直接走 R2；需要手动指定依赖镜像时可通过 `UV_INDEX_URL` 环境变量覆盖。 5. 在 3D View 按 `N`，打开 `SnapForm` 页签。 6. 如需其他模型，展开 `Models` panel，找到模型并点击该行的 `Download`。 diff --git a/docs/index.md b/docs/index.md index 83d164a..14072d4 100644 --- a/docs/index.md +++ b/docs/index.md @@ -7,4 +7,5 @@ - [Surface](surface.md)：透视网格、Geometry Nodes 后处理与材质法线 - [输出格式](outputs.md)：Debug 输出类型、参数和任务产物 - [AI Service](service.md)：Environment、模型缓存、任务目录与运行机制 +- [BlendJob 开发者指南](blendjob.md)：可复用 Job Server、Blender Client 与装饰器注册 API - [开发与许可证](development.md)：开发验证、打包、上游版本与许可证 diff --git a/docs/service.md b/docs/service.md index c02cfe4..98af95b 100644 --- a/docs/service.md +++ b/docs/service.md @@ -2,7 +2,9 @@ ## 运行架构 -SnapForm 在独立的 Python 3.12 Environment 中运行常驻 HTTP Service，Blender 只通过标准库 `urllib.request` 提交任务，不把 AI 运行库载入 Blender Python。Service 只监听 `127.0.0.1`，请求使用随机 Token；同一 Storage Root 由 `service.lock` 保证只有一个实例。 +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`： @@ -12,12 +14,30 @@ Environment、模型和任务结果默认保存在 `~/.snapform`： ├── models/ ├── jobs/ ├── manifest.json -├── service.log -├── service.json # Service 运行时存在 -└── service.lock # Service 运行时持有 +├── runtime.json # 持久的模型下载源状态 +├── environment.log +├── environment-status.json +└── service.log ``` -不同 Blender 版本可共享此目录。也可以在 Extension Preferences 的 `Storage Root` 中修改位置。 +不同 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 @@ -41,7 +61,9 @@ Windows 当前发行环境提供 DirectML 与 CPU；`Auto` 优先使用 `DmlExec - DA3 Small、DA3 Base - BEN2 Base -首次安装自动下载 `MoGe-2 ViT-S Normal` 与 FP16 `BEN2 Base`，其他模型在 `Models` panel 按需下载。模型按 Hugging Face 仓库末段名称保存在 `~/.snapform/models/`，并保留仓库内部结构。声明了文件清单与校验值的模型会直接走 Hugging Face `resolve` 直链流式下载，跳过文件清单与 ETag 预校验；连接源不可达时使用 3 秒 TCP 与 TLS 建连超时，随后回退到 `https://models.omoolab.xyz//`，并把 `models_from_r2` 写入 service.json，本次 service 周期内后续模型直接走 R2，不再重复等待 Hugging Face。建连成功后不限制模型文件读取时长。R2 文件会校验大小与 SHA-256。未声明清单的仓库仍使用 `huggingface_hub` 快照下载。可用 `HF_ENDPOINT` 指定 Hugging Face 镜像。 +首次安装由 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 后处理中恢复。 @@ -56,9 +78,15 @@ uv 默认从 Astral 的 GitHub Release 下载，TCP 与 TLS 建连超过 3 秒 ## 生命周期 -扩展启用后只查找已有 Service；第一次执行 AI 操作或点击 `Start` 时才启动。每个 Blender 客户端每 5 秒发送心跳，正常关闭时注销。最后一个客户端离开后，Service 保留模型缓存 20 秒再退出。 +扩展启用后不启动 Service；第一次执行 AI 操作或点击 `Start` 时才为当前 Blender 启动专属实例。正常关闭或禁用扩展时直接停止自己启动的进程；服务也监视 Blender 父进程 PID，父进程异常退出后会自行停止。 -Service 全局同时执行一个 AI 任务。BEN2 独立缓存，深度模型只保留最近使用的一个；切换模型时释放旧会话。取消任务不会重启 Service，ONNX Runtime 当前调用返回后在最近检查点结束。 +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`。 diff --git a/docs/surface.md b/docs/surface.md index 1be19e0..8b217ad 100644 --- a/docs/surface.md +++ b/docs/surface.md @@ -166,52 +166,21 @@ Image Empty 使用 Packed Image 时，SnapForm 会通过 Blender 写出临时 PN `DA3 Nested Giant Large 1.1`。Small 和 Base 没有 Gaussian head，无法由 DA3 生成 `gs_ply`。 -每个任务在 `~/.snapform/jobs/` 下用独立文件夹包裹。文件夹统一采用 -`动作-对象_日期-时间-微秒`，每种任务都有唯一名称，例如 -`generate-depth-plane_20260807-153000-123456`。微秒用于避免同一秒启动的任务共用目录。 +每个 Service Job 被接受时，BlendJob 在 `~/.snapform/jobs/` 下创建一个独立文件夹。文件夹名是 32 位 opaque Job ID，不包含任务类型或对象名称。执行逻辑从 `JobContext.directory` 获取路径，输出产物直接写入该目录。 -Blender 为 Packed Image 或序列帧生成的输入副本也归档到对应 job 根目录：单图命名为 -`input.png`，序列帧命名为 `frame.0001.*`。它们以前暂存在 `jobs/inputs/`，现在不再创建这个 -公共中转文件夹。普通磁盘文件和视频仍直接读取原文件，不额外复制。 +Blender 为 Packed Image 或序列帧生成的输入副本先保存在系统临时目录，Service 读取完成后由 Blender 清理，不归档到 Job 目录。普通磁盘文件和视频仍直接读取原文件。Service Job 日志统一写入 `service.log`，进度只保存在 Service 内存中。Environment Bootstrap 使用 Storage Root 下的 `environment.log` 与 `environment-status.json`，不创建 Job 目录。 -日志、运行中的状态文件、归档输入与输出产物均直接放在 job 根目录。各类任务名称为： - -- `install-environment` -- `download-model` -- `export-reconstruction` -- `generate-depth-plane` -- `generate-image-mesh` -- `remove-background` -- `remove-image-background` -- `generate-original-depth-plane` +与几何输出相关的 Job Type 为 `generate-da3-depth`、`generate-moge2-depth`、`generate-da3-surface`、`generate-moge2-surface`、`debug-da3` 与 `debug-moge2`。每种 Job 只接收对应模型与输出所需的参数。 DA3 ONNX exporter 产生的 `depth_vis/`、`exports/mini_npz/` 会在完成后自动平铺， 不会保留只含一层产物的子目录： ```text -export-reconstruction_20260807-153000-123456/ -├── job.log -├── status.json -└── scene.glb +/ +├── scene.glb、*.ply、*.jpg 或 results.npz +└── 其它对应 Debug 格式的产物 -export-reconstruction_20260807-153000-123456/ -├── job.log -├── status.json -└── *.ply - -export-reconstruction_20260807-153000-123456/ -├── job.log -├── status.json -└── *.jpg - -export-reconstruction_20260807-153000-123456/ -├── job.log -├── status.json -└── results.npz - -remove-background_20260807-153000-123456/ -├── job.log -├── status.json +/ ├── foreground.png 单图结果 ├── foreground.0001.png… 视频/序列帧结果 ├── frame.0001.jpg… 视频/序列帧的源帧副本 diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index 8a45708..2641f33 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.py @@ -1,7 +1,9 @@ import bpy from . import client -from .image_depth import ( +from .blender import depth_plane as image_depth +from .blender import surface as image_mesh +from .blender.depth_plane import ( ApplyRemovedBackground, BuildDepthPlane, BuildOriginalDepthPlane, @@ -11,11 +13,11 @@ from .image_depth import ( SnapFormImageMenu, draw_image_context_menu, ) -from .image_mesh import BuildImageMesh, ConvertImageToMesh +from .blender.surface import BuildImageMesh, ConvertImageToMesh from .operators import ( CancelJob, ControlService, - RunJob, + JOB_OPERATOR_CLASSES, cancel_active_job, close_active_job, ) @@ -35,7 +37,7 @@ from .status_bar import draw_status_bar CLASSES = ( SnapFormPreferences, SnapFormSettings, - RunJob, + *JOB_OPERATOR_CLASSES, CancelJob, ControlService, SnapFormImageMenu, diff --git a/src/snapform/ai_service/__init__.py b/src/snapform/ai_service/__init__.py new file mode 100644 index 0000000..540ce8f --- /dev/null +++ b/src/snapform/ai_service/__init__.py @@ -0,0 +1 @@ +"""SnapForm AI Service implementation without Blender dependencies.""" diff --git a/src/snapform/ben2_onnx.py b/src/snapform/ai_service/ben2_onnx.py similarity index 100% rename from src/snapform/ben2_onnx.py rename to src/snapform/ai_service/ben2_onnx.py diff --git a/src/snapform/da3_onnx.py b/src/snapform/ai_service/da3_onnx.py similarity index 100% rename from src/snapform/da3_onnx.py rename to src/snapform/ai_service/da3_onnx.py diff --git a/src/snapform/worker.py b/src/snapform/ai_service/jobs.py similarity index 98% rename from src/snapform/worker.py rename to src/snapform/ai_service/jobs.py index 1c886f3..b31f22c 100644 --- a/src/snapform/worker.py +++ b/src/snapform/ai_service/jobs.py @@ -84,7 +84,7 @@ def check_cancelled(cancel_check): def write_status( - path, + target, progress, message, error=None, @@ -111,7 +111,19 @@ def write_status( else: state = "running" payload["state"] = state - Path(path).write_text( + 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", ) @@ -991,15 +1003,15 @@ def progress_range(args, staged): return 0.05, 0.90 -def _discovery_path(model_dir): - """storage_root/service.json，从 storage_root/models/ 反推 storage_root。""" - return Path(model_dir).parent.parent / "service.json" +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( - _discovery_path(model_dir).read_text(encoding="utf-8") + _runtime_state_path(model_dir).read_text(encoding="utf-8") ) except (OSError, ValueError, json.JSONDecodeError): return False @@ -1007,8 +1019,8 @@ def _models_from_r2(model_dir): def _mark_models_from_r2(model_dir): - """在 service.json 上记录 models_from_r2，保留文件里已有的其它字段。""" - path = _discovery_path(model_dir) + """记录当前运行周期的模型镜像选择。""" + path = _runtime_state_path(model_dir) try: payload = json.loads(path.read_text(encoding="utf-8")) except (OSError, ValueError, json.JSONDecodeError): @@ -1082,6 +1094,9 @@ def _download_declared_or_fail(args, model_dir, label, start, span, stage_kwargs 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。 @@ -1139,6 +1154,9 @@ def _download_snapshot_or_fail(args, model_dir, label, start, span, stage_kwargs 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( @@ -1186,6 +1204,7 @@ def _download_file_from_url( 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"}) @@ -1199,6 +1218,8 @@ def _download_file_from_url( 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 @@ -1217,6 +1238,8 @@ def _download_file_from_url( temporary.replace(destination) return except Exception as error: + if cancel_check is not None: + cancel_check() last_error = error try: temporary.unlink() @@ -1282,6 +1305,7 @@ def _download_declared_files( checksum, progress=on_chunk, retries=retries, + cancel_check=getattr(args, "cancel_check", None), ) diff --git a/src/snapform/ai_service/model_manager.py b/src/snapform/ai_service/model_manager.py new file mode 100644 index 0000000..c21f17a --- /dev/null +++ b/src/snapform/ai_service/model_manager.py @@ -0,0 +1,153 @@ +import gc +import time +from pathlib import Path +from types import SimpleNamespace + +import jobs +import models as model_catalog + + +class ModelSessionCache: + """Cache ONNX sessions shared by jobs in one service process.""" + + def __init__(self): + self.ben2_model = None + self.ben2_key = None + self.da3_model = None + self.da3_key = None + self.moge_model = None + self.moge_key = None + + def get_ben2(self, model_dir, device): + requested_device = getattr(device, "type", device) + key = (str(Path(model_dir).resolve()), requested_device, "onnx") + if self.ben2_model is not None and self.ben2_key == key: + return self.ben2_model, 0.0 + self.release_ben2() + return self._load_ben2(model_dir, requested_device, key) + + def _load_ben2(self, model_dir, requested_device, key): + from ben2_onnx import create_session + + started = time.perf_counter() + model = create_session(model_dir, requested_device) + self.ben2_model = model + self.ben2_key = key + elapsed_ms = (time.perf_counter() - started) * 1000.0 + return model, elapsed_ms + + def get_da3(self, model_dir, device): + requested_device = getattr(device, "type", device) + key = (str(Path(model_dir).resolve()), requested_device, "onnx") + if self.moge_model is not None: + self.release_moge() + if self.da3_model is not None and self.da3_key == key: + return self.da3_model + self.release_da3() + return self._load_da3(model_dir, requested_device, key) + + def _load_da3(self, model_dir, requested_device, key): + from da3_onnx 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 + + key = (str(Path(model_dir).resolve()), requested_device, "onnx") + if self.da3_model is not None: + self.release_da3() + if self.moge_model is not None and self.moge_key == key: + return self.moge_model + self.release_moge() + model = create_session(model_dir, requested_device) + self.moge_model = model + self.moge_key = key + return model + + def release_ben2(self): + self.ben2_model = None + self.ben2_key = None + self._release_runtime_cache() + + def release_da3(self): + self.da3_model = None + self.da3_key = None + self._release_runtime_cache() + + def release_moge(self): + self.moge_model = None + self.moge_key = None + self._release_runtime_cache() + + def close(self): + self.ben2_model = None + self.da3_model = None + self.moge_model = None + self.ben2_key = None + self.da3_key = None + self.moge_key = None + self._release_runtime_cache() + + def snapshot(self): + return { + "ben2": Path(self.ben2_key[0]).name if self.ben2_key else "", + "da3": Path(self.da3_key[0]).name if self.da3_key else "", + "moge2": Path(self.moge_key[0]).name if self.moge_key else "", + } + + def _release_runtime_cache(self): + gc.collect() + + +class ModelManager: + """Own downloaded model files and loaded inference sessions.""" + + def __init__(self, models_directory): + self.models_directory = Path(models_directory) + self.sessions = ModelSessionCache() + + def directory(self, model_key): + model = model_catalog.get_downloadable_model(model_key) + return self.models_directory / model.directory_name + + def get_ben2(self, model_dir, device): + return self.sessions.get_ben2(model_dir, device) + + 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 snapshot(self): + return {"loaded": self.sessions.snapshot()} + + def clear(self): + self.sessions.close() + + def close(self): + self.sessions.close() + + def download(self, context, model_key): + model = model_catalog.get_downloadable_model(model_key) + args = SimpleNamespace( + repository=model.repository, + model_dir=str(self.directory(model_key)), + status=context.update, + cancel_check=context.check_cancelled, + label=model.label, + allow_pattern=list(model.download_patterns), + required_pattern=list(model.ready_patterns), + r2_directory=model.r2_directory, + r2_file=[ + f"{filename}|{size}|{checksum}" + for filename, size, checksum in model.r2_files + ], + ) + context.check_cancelled() + jobs.download_model(args) + context.check_cancelled() diff --git a/src/snapform/moge_onnx.py b/src/snapform/ai_service/moge_onnx.py similarity index 100% rename from src/snapform/moge_onnx.py rename to src/snapform/ai_service/moge_onnx.py diff --git a/src/snapform/onnx_runtime.py b/src/snapform/ai_service/onnx_runtime.py similarity index 100% rename from src/snapform/onnx_runtime.py rename to src/snapform/ai_service/onnx_runtime.py diff --git a/src/snapform/ai_service/server.py b/src/snapform/ai_service/server.py new file mode 100644 index 0000000..c419f61 --- /dev/null +++ b/src/snapform/ai_service/server.py @@ -0,0 +1,145 @@ +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/__init__.py b/src/snapform/blender/__init__.py new file mode 100644 index 0000000..c741b19 --- /dev/null +++ b/src/snapform/blender/__init__.py @@ -0,0 +1 @@ +"""Blender-process geometry construction and import logic.""" diff --git a/src/snapform/image_depth.py b/src/snapform/blender/depth_plane.py similarity index 98% rename from src/snapform/image_depth.py rename to src/snapform/blender/depth_plane.py index 2ddb810..1611eb2 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/blender/depth_plane.py @@ -53,7 +53,7 @@ VIDEO_NAME_SUFFIXES = ( def image_depth_model_items(owner, context): - from .models import depth_plane_model_enum_items + from ..models import depth_plane_model_enum_items return depth_plane_model_enum_items(owner, context) @@ -228,7 +228,7 @@ def draw_process_resolution(layout, owner, image): def save_packed_image(image, scene): - from .client import temporary_input_path + from ..client import temporary_input_path output_path = temporary_input_path() image_settings = scene.render.image_settings @@ -286,7 +286,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 + from ..client import temporary_input_directory temporary_directory = Path(temporary_input_directory()) try: @@ -1487,7 +1487,7 @@ def set_displacement_only(material): def node_asset_path(): - return Path(__file__).with_name("assets") / NODE_ASSET_NAME + return Path(__file__).parent.parent / "assets" / NODE_ASSET_NAME def load_depth_plane_node_group(group_name=DEPTH_PLANE_NODE_GROUP_NAME): @@ -2227,9 +2227,9 @@ 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 ..operators import active_job + from ..client import service_busy, model_ready + from ..models import BEN2_MODEL return ( is_image_empty(context.object) @@ -2239,8 +2239,6 @@ class RemoveImageBackground(bpy.types.Operator): ) def execute(self, context): - from .operators import IMAGE_BACKGROUND_JOB - source_object = context.object total_frames = video_total_frames(source_object.data) if total_frames > 1: @@ -2256,8 +2254,7 @@ class RemoveImageBackground(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} - result = bpy.ops.snapform.run_job( - job_type=IMAGE_BACKGROUND_JOB, + 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, @@ -2312,8 +2309,8 @@ class ConvertImageToDepthPlane(bpy.types.Operator): @classmethod def poll(cls, context): - from .operators import active_job - from .client import service_busy + from ..operators import active_job + from ..client import service_busy return ( is_image_empty(context.object) @@ -2337,13 +2334,6 @@ class ConvertImageToDepthPlane(bpy.types.Operator): layout.prop(self, "snapform_keep_original") def execute(self, context): - from .operators import ( - IMAGE_DEPTH_JOB, - MOGE2_DEPTH_JOB, - ORIGINAL_BACKGROUND_JOB, - ORIGINAL_FRAMES_JOB, - ) - source_object = context.object if video_total_frames(source_object.data) > 1: context.scene.snapform_settings.video_frames = ( @@ -2358,12 +2348,12 @@ class ConvertImageToDepthPlane(bpy.types.Operator): except RuntimeError as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} - result = bpy.ops.snapform.run_job( - job_type=( - ORIGINAL_BACKGROUND_JOB - if self.snapform_remove_background - else ORIGINAL_FRAMES_JOB - ), + run_original = ( + bpy.ops.snapform.generate_original_background + if self.snapform_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, @@ -2385,8 +2375,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): except RuntimeError as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} - result = bpy.ops.snapform.run_job( - job_type=MOGE2_DEPTH_JOB, + result = bpy.ops.snapform.generate_moge2_depth( snapform_input_path=str(source_path), snapform_source_object_name=source_object.name, snapform_moge_resolution_level=( @@ -2422,8 +2411,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} - result = bpy.ops.snapform.run_job( - job_type=IMAGE_DEPTH_JOB, + 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), @@ -2447,17 +2435,16 @@ class SnapFormImageMenu(bpy.types.Menu): bl_label = "SnapForm" def draw(self, _context): - from .client import env_installed + from ..client import env_installed layout = self.layout layout.operator_context = "INVOKE_DEFAULT" if not env_installed(): - operator = layout.operator( - "snapform.run_job", + layout.operator( + "snapform.install_environment", text="Install Environment", icon="IMPORT", ) - operator.job_type = "INSTALL_ENV" return layout.operator( "snapform.convert_image_to_depth_plane", diff --git a/src/snapform/image_mesh.py b/src/snapform/blender/surface.py similarity index 94% rename from src/snapform/image_mesh.py rename to src/snapform/blender/surface.py index cc3fff9..e9f3844 100644 --- a/src/snapform/image_mesh.py +++ b/src/snapform/blender/surface.py @@ -2,7 +2,7 @@ from pathlib import Path import bpy -from .image_depth import ( +from .depth_plane import ( MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME, create_material, draw_process_resolution, @@ -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 ..models import surface_model_enum_items return surface_model_enum_items(owner, context) @@ -329,7 +329,7 @@ def depth_discontinuity_ratio(depth): def node_asset_path(): - return Path(__file__).with_name("assets") / NODE_ASSET_NAME + return Path(__file__).parent.parent / "assets" / NODE_ASSET_NAME def load_geometry_node_group(group_name=NODE_GROUP_NAME): @@ -502,7 +502,7 @@ def build_image_mesh( depth_scale = material_depth_scale(mesh_object, camera_parameters) normal_image = None if moge2 and attributes.get("normal_map", np.empty(0)).size: - from .image_depth import create_moge_normal_image + from .depth_plane import create_moge_normal_image normal_map = attributes["normal_map"].copy() normal_map[..., 0] /= camera_parameters["scale_x"] @@ -651,9 +651,9 @@ 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 ..models import MODELS, MOGE2_MODELS + from ..operators import active_job + from ..client import service_busy, model_ready return ( is_image_empty(context.object) @@ -683,8 +683,6 @@ class ConvertImageToMesh(bpy.types.Operator): layout.prop(self, "snapform_keep_original") def execute(self, context): - from .operators import IMAGE_MESH_JOB, MOGE2_MESH_JOB - source_object = context.object if video_total_frames(source_object.data) != 1: self.report({"ERROR"}, "Convert to Surface currently supports still images") @@ -703,24 +701,29 @@ class ConvertImageToMesh(bpy.types.Operator): except (RuntimeError, ValueError) as error: self.report({"ERROR"}, str(error)) return {"CANCELLED"} - result = bpy.ops.snapform.run_job( - job_type=( - MOGE2_MESH_JOB - if self.snapform_model.startswith("MOGE2_") - else IMAGE_MESH_JOB - ), - snapform_input_path=str(source_path), - snapform_source_object_name=source_object.name, - snapform_model_key=self.snapform_model, - snapform_use_ray_pose=self.snapform_use_ray_pose, - snapform_process_res=process_res, - snapform_moge_resolution_level=self.snapform_moge_resolution_level, - 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, - ) + 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, + } + if self.snapform_model.startswith("MOGE2_"): + result = bpy.ops.snapform.generate_moge2_surface( + **common, + snapform_moge_resolution_level=( + self.snapform_moge_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, + ) if "CANCELLED" in result: if delete_input: remove_temporary_input(source_path) diff --git a/src/snapform/blendjob/__init__.py b/src/snapform/blendjob/__init__.py new file mode 100644 index 0000000..914b8d9 --- /dev/null +++ b/src/snapform/blendjob/__init__.py @@ -0,0 +1,37 @@ +"""Reusable local HTTP job service for Blender add-ons.""" + +from .client import JobClient, JobRequest +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", + "JobServer", + "ServiceConnection", + "ServiceController", +) + +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 new file mode 100644 index 0000000..d407ff1 --- /dev/null +++ b/src/snapform/blendjob/client.py @@ -0,0 +1,99 @@ +import json +from dataclasses import dataclass +from urllib.error import HTTPError, URLError +from urllib.request import ProxyHandler, Request, build_opener + + +class JobClient: + """Dependency-free HTTP client for a local :class:`JobServer`.""" + + def __init__(self, host, port, *, timeout=2.0): + self.base_url = f"http://{host}:{int(port)}" + self.timeout = timeout + self.opener = build_opener(ProxyHandler({})) + + def request(self, method, path, payload=None, timeout=None): + data = None + headers = {} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + request = Request( + f"{self.base_url}{path}", + data=data, + headers=headers, + method=method, + ) + try: + with self.opener.open( + request, + timeout=self.timeout if timeout is None else timeout, + ) as response: + return json.loads(response.read().decode("utf-8")) + except HTTPError as error: + try: + detail = json.loads(error.read().decode("utf-8")).get("detail") + except (json.JSONDecodeError, UnicodeDecodeError): + detail = None + raise RuntimeError(detail or f"Job Server returned HTTP {error.code}") + except (OSError, URLError) as error: + raise RuntimeError(f"Unable to reach Job Server: {error}") + + def health(self): + return self.request("GET", "/health") + + def submit(self, job_type, **parameters): + return self.request( + "POST", + "/jobs", + { + "job_type": job_type, + "parameters": 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}") + + def cancel(self, job_id): + return self.request("DELETE", f"/jobs/{job_id}") + + def resources(self): + return self.request("GET", "/resources") + + def resource(self, name): + return self.request("GET", f"/resources/{name}") + + def clear_resource(self, name): + return self.request("POST", f"/resources/{name}/clear") + + def shutdown(self): + return self.request("POST", "/shutdown") +@dataclass(frozen=True, init=False) +class JobRequest: + """Serializable request created before a Job is submitted.""" + + 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 {})) + + @classmethod + def create(cls, job_type, **parameters): + return cls(job_type, parameters) + + @property + def payload(self): + return { + "job_type": self.job_type, + "parameters": self.parameters, + } diff --git a/src/snapform/blendjob/operator.py b/src/snapform/blendjob/operator.py new file mode 100644 index 0000000..65bfb24 --- /dev/null +++ b/src/snapform/blendjob/operator.py @@ -0,0 +1,267 @@ +import threading +from dataclasses import dataclass +from pathlib import Path + +import bpy + +from .client import JobRequest + + +@dataclass +class ModalJobState: + """Runtime state shared by all Blender modal Job operators.""" + + controller: object + data: object = None + cleanup: object = None + job_id: str = "" + directory: Path | None = None + started: bool = False + cancelled: bool = False + start_error: str = "" + start_thread: threading.Thread | None = None + 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) + + +_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 + try: + job.controller.cancel(job.job_id) + except RuntimeError: + pass + + +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.""" + + bl_options = {"INTERNAL"} + poll_interval = 0.25 + starting_message = "Starting task" + controller_factory = 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 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 job_parameters(self, _context): + """Return JSON-compatible business parameters for this invocation.""" + return {} + + def create_job_data(self, _context): + """Return optional plug-in data retained until completion.""" + return None + + def execute(self, context): + global _active_job + + if active_job() is not None: + self.report({"WARNING"}, "Another task is running") + 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) + 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 + 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)) + 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) + job.start_thread = threading.Thread( + target=self._submit_job, + args=(job, request), + name="BlendJobSubmit", + daemon=True, + ) + job.start_thread.start() + return {"RUNNING_MODAL"} + + def _submit_job(self, job, request): + if job.cancelled: + job.started = True + return + try: + submitted = job.controller.submit(request) + job.job_id = submitted["job_id"] + job.directory = Path(submitted["directory"]) + job.started = True + if job.cancelled: + job.controller.cancel(job.job_id) + except (KeyError, OSError, RuntimeError, TypeError) as error: + job.start_error = str(error) + + def modal(self, context, event): + if event.type == "ESC": + cancel_active_job() + if event.type != "TIMER": + return {"PASS_THROUGH"} + + job = active_job() + if job is None: + self.remove_job_timer(context) + return {"CANCELLED"} + if job.start_error: + return self._finish_start_error(context, job) + if not job.started: + return {"RUNNING_MODAL"} + + try: + status = job.controller.status(job.job_id) + except RuntimeError as error: + return self._finish_failure(context, job, str(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) + 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: + return self._finish_failure(context, job, error) + return self._close( + context, + job, + message or status.get("message") or "Task complete", + cancelled=False, + ) + + 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) + 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") + if stage is not None and stage != job.progress_stage: + job.progress_stage = stage + job.progress = reported + 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 + + def format_job_status(self, status, fallback): + message = status.get("message") + stage = status.get("stage") + stages = status.get("stages") + stage_label = status.get("stage_label") + if stage and stages: + prefix = f"{stage}/{stages}" + normalized = str(message or "").strip() + if normalized and normalized.lower() != str(stage_label or "").lower(): + return f"{prefix} {normalized}" + return f"{prefix} {stage_label}" if stage_label else prefix + return str(message or stage_label or fallback).strip() + + def cancel(self, context): + cancel_active_job() + self.remove_job_timer(context) + self.redraw_job_ui(context, force=True) diff --git a/src/snapform/blendjob/server.py b/src/snapform/blendjob/server.py new file mode 100644 index 0000000..b7d8bb3 --- /dev/null +++ b/src/snapform/blendjob/server.py @@ -0,0 +1,350 @@ +import os +import threading +import time +import traceback +import uuid +from contextlib import asynccontextmanager +from pathlib import Path + + +TERMINAL_STATES = frozenset({"succeeded", "failed", "cancelled"}) + + +class JobCancelled(RuntimeError): + """Raised by a handler when its job has been cancelled.""" + + +class JobContext: + """Mutable state and cancellation API passed to a registered job handler.""" + + def __init__(self, job_id, job_type, parameters, resources, directory): + self.job_id = job_id + self.job_type = job_type + self.parameters = parameters + self.directory = Path(directory) + self._resources = resources + self._lock = threading.Lock() + self._cancel_requested = False + self._status = { + "job_id": job_id, + "job_type": job_type, + "state": "queued", + "progress": 0.0, + "message": "Task queued", + "directory": str(self.directory), + } + + def update(self, progress, message, error=None, state=None, **details): + """Publish a progress snapshot returned by ``GET /jobs/{job_id}``.""" + if state is None: + if error: + state = "failed" + elif progress >= 1.0: + state = "succeeded" + else: + state = "running" + snapshot = { + "job_id": self.job_id, + "job_type": self.job_type, + "state": state, + "progress": min(max(float(progress), 0.0), 1.0), + "message": str(message), + "directory": str(self.directory), + **{key: value for key, value in details.items() if value is not None}, + } + if error: + snapshot["error"] = str(error) + with self._lock: + self._status = snapshot + + def snapshot(self): + with self._lock: + return dict(self._status) + + def request_cancel(self): + with self._lock: + if self._status["state"] in TERMINAL_STATES: + return False + self._cancel_requested = True + self._status = { + **self._status, + "state": "cancelling", + "message": "Cancelling task", + } + return True + + def is_cancelled(self): + with self._lock: + return self._cancel_requested + + def check_cancelled(self): + if self.is_cancelled(): + raise JobCancelled("Task cancelled") + + def resource(self, name): + """Return a shared resource registered on the owning JobServer.""" + try: + return self._resources[name] + except KeyError: + raise KeyError(f"Unknown service 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): + self.name = name + self.max_job_history = max(int(max_job_history), 0) + self.handlers = {} + self.resources = {} + self.jobs = {} + self.active_job_id = None + self.lock = threading.Lock() + self.shutdown_event = threading.Event() + self.started_at = time.time() + self.jobs_directory = Path(jobs_directory or Path.cwd() / "jobs") + + def add_resource(self, name, resource): + """Register shared state owned for the full service lifetime.""" + if not name or "/" in name: + raise ValueError(f"Invalid resource name: {name}") + with self.lock: + if name in self.resources: + raise ValueError(f"Resource is already registered: {name}") + self.resources[name] = resource + return resource + + def resource_snapshot(self, name): + try: + resource = self.resources[name] + except KeyError: + raise KeyError(name) from None + snapshot = getattr(resource, "snapshot", None) + return snapshot() if snapshot is not None else {} + + def resource_snapshots(self): + return { + name: self.resource_snapshot(name) + for name in tuple(self.resources) + } + + def clear_resource(self, name): + """Clear an idle service resource without racing a new job submission.""" + with self.lock: + if self.active_job_id is not None: + return False + try: + resource = self.resources[name] + except KeyError: + raise KeyError(name) from None + clear = getattr(resource, "clear", None) + if clear is None: + raise TypeError(f"Resource cannot be cleared: {name}") + clear() + return True + + def close(self): + for resource in reversed(tuple(self.resources.values())): + close = getattr(resource, "close", None) + if close is not None: + close() + + def job(self, job_type): + """Register ``handler(context, parameters)`` for a public job type.""" + def register(handler): + if job_type in self.handlers: + raise ValueError(f"Job type is already registered: {job_type}") + self.handlers[job_type] = handler + return handler + + return register + + def snapshot(self, instance_id): + with self.lock: + active = self.jobs.get(self.active_job_id) + result = { + "ready": True, + "service": self.name, + "instance_id": instance_id, + "busy": active is not None, + "active_job": active.snapshot() if active else None, + "started_at": self.started_at, + "resources": self.resource_snapshots(), + } + return result + + def submit(self, job_type, parameters, *, job_id=None): + 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) + directory = self.jobs_directory / job_id + try: + directory.mkdir() + except FileExistsError: + raise ValueError(f"Job directory already exists: {job_id}") from None + context = JobContext( + job_id, + job_type, + parameters, + self.resources, + directory, + ) + self.jobs[job_id] = context + self.active_job_id = job_id + self._prune_jobs() + return context + + def _prune_jobs(self): + completed = [ + identifier + for identifier, context in self.jobs.items() + if identifier != self.active_job_id + and context.snapshot()["state"] in TERMINAL_STATES + ] + excess = len(completed) - self.max_job_history + for identifier in completed[:max(excess, 0)]: + self.jobs.pop(identifier, None) + + def run(self, context): + context.update(0.0, "Service accepted the task", state="running") + try: + context.check_cancelled() + 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") + except JobCancelled: + 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") + finally: + with self.lock: + if self.active_job_id == context.job_id: + self.active_job_id = None + + def create_app(self, instance_id): + from fastapi import BackgroundTasks, FastAPI, HTTPException + + @asynccontextmanager + async def lifespan(_app): + yield + self.close() + + app = FastAPI( + title=self.name, + docs_url=None, + redoc_url=None, + openapi_url=None, + lifespan=lifespan, + ) + app.state.job_server = self + + @app.get("/health") + def health(): + return self.snapshot(instance_id) + + @app.get("/resources") + def resources(): + return self.resource_snapshots() + + @app.get("/resources/{resource_name}") + def resource_status( + resource_name: str, + ): + try: + return self.resource_snapshot(resource_name) + except KeyError: + raise HTTPException(status_code=404, detail="Resource was not found") + + @app.post("/resources/{resource_name}/clear") + def clear_resource( + resource_name: str, + ): + try: + cleared = self.clear_resource(resource_name) + except KeyError: + raise HTTPException(status_code=404, detail="Resource was not found") + except TypeError as error: + raise HTTPException(status_code=405, detail=str(error)) + if not cleared: + raise HTTPException(status_code=409, detail="Service is busy") + return {"resource": resource_name, "state": "cleared"} + + @app.post("/jobs", status_code=202) + def submit_job( + payload: dict, + background_tasks: BackgroundTasks, + ): + job_type = str(payload.get("job_type", payload.get("command", ""))) + parameters = payload.get("parameters") + if not job_type or not isinstance(parameters, dict): + raise HTTPException(status_code=422, detail="Invalid job request") + try: + context = self.submit(job_type, parameters) + except KeyError: + 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}") + def job_status(job_id: str): + context = self.jobs.get(job_id) + if context is None: + raise HTTPException(status_code=404, detail="Job was not found") + return context.snapshot() + + @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(): + raise HTTPException(status_code=404, detail="Active job was not found") + return context.snapshot() + + @app.post("/shutdown", status_code=202) + def shutdown(): + self.shutdown_event.set() + return {"state": "stopping"} + + return app + + +def watch_parent(parent_pid, shutdown_event, interval=1.0): + """Stop a dedicated service after its owning Blender process exits.""" + if not parent_pid: + return + while not shutdown_event.wait(interval): + if not _process_exists(parent_pid): + shutdown_event.set() + return + + +def _process_exists(process_id): + if os.name != "nt": + try: + os.kill(process_id, 0) + return True + except OSError: + return False + import ctypes + + synchronize = 0x00100000 + handle = ctypes.windll.kernel32.OpenProcess(synchronize, False, process_id) + if not handle: + return False + try: + return ctypes.windll.kernel32.WaitForSingleObject(handle, 0) == 0x00000102 + finally: + ctypes.windll.kernel32.CloseHandle(handle) diff --git a/src/snapform/blendjob/service.py b/src/snapform/blendjob/service.py new file mode 100644 index 0000000..dd22ddf --- /dev/null +++ b/src/snapform/blendjob/service.py @@ -0,0 +1,208 @@ +import socket +import subprocess +import time +import uuid +from dataclasses import dataclass + +from .client import JobClient + + +@dataclass +class ServiceConnection: + port: int + pid: int + instance_id: str + client: JobClient + + +class ServiceController: + """Own a dedicated local JobServer process and its HTTP connection.""" + + def __init__( + self, + command_factory, + *, + cwd_factory, + log_path_factory, + host="127.0.0.1", + creationflags=0, + ): + self.command_factory = command_factory + self.cwd_factory = cwd_factory + self.log_path_factory = log_path_factory + self.host = host + self.creationflags = creationflags + self.connection = None + self.process = None + self.log_file = None + self.auto_start = False + self.snapshot = {"state": "STOPPED", "message": "Stopped"} + + def start(self): + if self.process is not None and self.process.poll() is None: + return self.connection + self._close_process() + port = self._available_port() + instance_id = uuid.uuid4().hex + log_path = self.log_path_factory() + log_path.parent.mkdir(parents=True, exist_ok=True) + self.log_file = log_path.open("a", encoding="utf-8") + command = self.command_factory(port, instance_id) + try: + self.process = subprocess.Popen( + command, + stdout=self.log_file, + stderr=subprocess.STDOUT, + cwd=self.cwd_factory(), + creationflags=self.creationflags, + ) + except OSError: + self.log_file.close() + self.log_file = None + raise + client = JobClient(self.host, port) + self.connection = ServiceConnection( + port, + self.process.pid, + instance_id, + client, + ) + self.snapshot = {"state": "STARTING", "message": "Starting", "port": port} + return self.connection + + def health(self, timeout=0.3): + if self.connection is None: + raise RuntimeError("Service 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") + state = "BUSY" if health.get("busy") else "READY" + self.snapshot = { + **health, + "state": state, + "message": "Busy" if state == "BUSY" else "Ready", + "port": self.connection.port, + } + return health + + def ensure(self, timeout=60.0): + self.auto_start = True + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + try: + self.health(timeout=0.5) + return self.connection + except RuntimeError: + self._close_finished_process() + if self.process is None: + self.start() + time.sleep(0.05) + raise RuntimeError(f"Service did not become ready; log: {self.log_path_factory()}") + + def enable(self): + self.auto_start = True + try: + self.health(timeout=0.2) + except RuntimeError: + self.start() + return self.connection + + def stop(self, timeout=10.0): + self.auto_start = False + if self.connection is not None: + try: + self.connection.client.shutdown() + except RuntimeError: + pass + deadline = time.monotonic() + timeout + while self.process is not None and self.process.poll() is None: + if time.monotonic() >= deadline: + self._close_process(terminate=True) + break + time.sleep(0.05) + self._close_process() + self.connection = None + self.snapshot = {"state": "STOPPED", "message": "Stopped"} + + def restart(self): + self.stop() + self.auto_start = True + return self.start() + + def submit(self, request): + return self.ensure().client.submit_request(request) + + def status(self, job_id): + return self.ensure().client.status(job_id) + + def cancel(self, job_id): + return self.ensure().client.cancel(job_id) + + def resource(self, name): + return self.ensure().client.resource(name) + + def clear_resource(self, name): + return self.ensure().client.clear_resource(name) + + def detach(self): + self.auto_start = False + self.connection = None + self._close_process(terminate=True) + self.snapshot = {"state": "STOPPED", "message": "Stopped"} + + def poll(self, *, available=True): + if not available: + self.snapshot = {"state": "UNAVAILABLE", "message": "Environment is not installed"} + return 2.0 + try: + self.health(timeout=0.5) + return 5.0 + except RuntimeError as error: + self._close_finished_process() + if not self.auto_start: + self.snapshot = {"state": "STOPPED", "message": "Stopped"} + return 2.0 + try: + self.start() + except OSError as start_error: + self.snapshot = {"state": "ERROR", "message": str(start_error)} + return 2.0 + self.snapshot = {"state": "STARTING", "message": "Starting"} + return 1.0 + + def mark_job_complete(self, job_id): + active = self.snapshot.get("active_job") + if active and active.get("job_id") != job_id: + return + if self.snapshot.get("state") == "BUSY": + self.snapshot = { + **self.snapshot, + "busy": False, + "active_job": None, + "state": "READY", + "message": "Ready", + } + + def _available_port(self): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind((self.host, 0)) + return listener.getsockname()[1] + + def _close_finished_process(self): + if self.process is not None and self.process.poll() is not None: + self._close_process() + self.connection = None + + def _close_process(self, terminate=False): + process = self.process + self.process = None + if process is not None and terminate and process.poll() is None: + process.terminate() + try: + process.wait(timeout=3.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=3.0) + if self.log_file is not None: + self.log_file.close() + self.log_file = None diff --git a/src/snapform/client.py b/src/snapform/client.py index fa00baa..ef47d67 100644 --- a/src/snapform/client.py +++ b/src/snapform/client.py @@ -1,26 +1,19 @@ import json import os import platform -import secrets -import shutil -import socket import subprocess import sys import tempfile import threading -import time -import uuid -from dataclasses import dataclass -from datetime import datetime from pathlib import Path -from urllib.error import HTTPError, URLError -from urllib.request import ProxyHandler, Request, build_opener 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, @@ -29,49 +22,60 @@ from .models import ( SERVICE_ADDRESS = "127.0.0.1" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" -JOB_FOLDER_NAMES = { - "INSTALL_ENV": "install-environment", - "DOWNLOAD_MODEL": "download-model", - "RECONSTRUCT": "export-reconstruction", - "IMAGE_DEPTH": "generate-depth-plane", - "MOGE2_DEPTH": "generate-moge2-depth-plane", - "MOGE2_MESH": "generate-moge2-surface", - "MOGE2_DEBUG": "test-moge2-api", - "IMAGE_MESH": "generate-image-mesh", - "REMOVE_BACKGROUND": "remove-background", - "IMAGE_BACKGROUND": "remove-image-background", - "ORIGINAL_BACKGROUND": "generate-original-depth-plane", - "ORIGINAL_FRAMES": "extract-original-frames", -} +_service_controller = None +_environment_controller = None -@dataclass(frozen=True) -class JobRequest: - payload: dict +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 -@dataclass -class Connection: - port: int - token: str - pid: int - instance_id: str - - -@dataclass -class ServiceProcess: - process: subprocess.Popen - log_file: object - - -_service_connection = None -_launched_service = None -_registered_instance_id = None -_client_id = uuid.uuid4().hex -_service_auto_start = False -_service_snapshot = {"state": "STOPPED", "message": "Stopped"} -_local_opener = build_opener(ProxyHandler({})) -_client_label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip() +def environment_controller(): + global _environment_controller + if _environment_controller is None: + _environment_controller = EnvironmentController() + return _environment_controller def preferences(): @@ -207,12 +211,6 @@ def production_device(): return configured or "AUTO" -def jobs_directory(): - path = storage_root() / "jobs" - path.mkdir(parents=True, exist_ok=True) - return path - - def temporary_input_path(): directory = Path(tempfile.mkdtemp(prefix="snapform-input-")) return directory / "input.png" @@ -222,52 +220,17 @@ def temporary_input_directory(): return Path(tempfile.mkdtemp(prefix="snapform-input-")) -def move_input_into_job(input_path, job_directory): - """Move a Blender-generated input copy into the flat job directory.""" - input_path = Path(input_path) - job_directory = Path(job_directory) - if input_path.is_file(): - destination = job_directory / "input.png" - input_path.replace(destination) - try: - input_path.parent.rmdir() - except OSError: - pass - return destination - if not input_path.is_dir(): - raise ValueError(f"Temporary input does not exist: {input_path}") - for source in input_path.iterdir(): - if not source.is_file(): - raise ValueError(f"Temporary input contains a directory: {source}") - destination = job_directory / source.name - if destination.exists(): - raise ValueError(f"Job input filename already exists: {source.name}") - source.replace(destination) - input_path.rmdir() - return job_directory - - -def job_folder_name(job_type): - """Return the folder name segment for a job type.""" - return JOB_FOLDER_NAMES.get(job_type, job_type.lower()) - - -def new_job_paths(job_type): - timestamp = datetime.now().strftime("%Y%m%d-%H%M%S-%f") - stem = f"{job_folder_name(job_type)}_{timestamp}" - directory = jobs_directory() / stem - directory.mkdir(parents=True, exist_ok=True) +def install_paths(): + """Return shared bootstrap paths outside the BlendJob workspace.""" + root = storage_root() return { - "log": directory / "job.log", - "status": directory / "status.json", - "result": directory, + "job_id": "", + "log": root / "environment.log", + "status": root / "environment-status.json", + "result": root, } -def job_id_from_status(status_path): - return Path(status_path).parent.name - - def uv_executable(): """Return the downloaded uv path inside the shared storage.""" return str(uv_cache_path()) @@ -303,282 +266,63 @@ def service_log_path(): return storage_root() / "service.log" -def service_discovery_path(): - return storage_root() / "service.json" - - -def service_lock_path(): - return storage_root() / "service.lock" - - -def _available_service_port(): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: - listener.bind((SERVICE_ADDRESS, 0)) - return listener.getsockname()[1] - - -def _service_url(connection, path): - return f"http://{SERVICE_ADDRESS}:{connection.port}{path}" - - -def _service_http_request(connection, method, path, payload=None, timeout=2.0): - data = None - headers = {"X-SnapForm-Token": connection.token} - if payload is not None: - data = json.dumps(payload).encode("utf-8") - headers["Content-Type"] = "application/json" - request = Request( - _service_url(connection, path), - data=data, - headers=headers, - method=method, - ) - try: - with _local_opener.open(request, timeout=timeout) as response: - return json.loads(response.read().decode("utf-8")) - except HTTPError as error: - try: - detail = json.loads(error.read().decode("utf-8")).get("detail") - except (json.JSONDecodeError, UnicodeDecodeError): - detail = None - raise RuntimeError( - detail or f"Service returned HTTP {error.code}" - ) - except (OSError, URLError) as error: - raise RuntimeError(f"Unable to reach Service: {error}") - - -def _read_service_discovery(): - try: - payload = json.loads(service_discovery_path().read_text(encoding="utf-8")) - return Connection( - port=int(payload["port"]), - token=str(payload["token"]), - pid=int(payload["pid"]), - instance_id=str(payload["instance_id"]), - ) - except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError): - return None - - -def _check_service(connection, timeout=0.3): - health = _service_http_request( - connection, - "GET", - "/health", - timeout=timeout, - ) - if health.get("instance_id") != connection.instance_id: - raise RuntimeError("Service discovery is stale") - return health - - -def _discover_live_service(timeout=0.3): - connection = _read_service_discovery() - if connection is None: - return None, None - try: - return connection, _check_service(connection, timeout=timeout) - except RuntimeError: - return None, None - - -def start_service(): - global _launched_service - - if _launched_service is not None and _launched_service.process.poll() is None: - return - _close_launched_service() - port = _available_service_port() - token = secrets.token_urlsafe(32) - instance_id = uuid.uuid4().hex - log_file = service_log_path().open("a", encoding="utf-8") - command = [ +def _service_command(port, instance_id): + return [ str(venv_python()), "-u", - str(Path(__file__).with_name("service.py")), + str(Path(__file__).with_name("ai_service") / "server.py"), "--host", SERVICE_ADDRESS, "--port", str(port), - "--token", - token, "--instance-id", instance_id, - "--discovery", - str(service_discovery_path()), - "--lock", - str(service_lock_path()), + "--parent-pid", + str(os.getpid()), ] - try: - process = subprocess.Popen( - command, - stdout=log_file, - stderr=subprocess.STDOUT, - cwd=storage_root(), + + +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(), ) - except OSError: - log_file.close() - raise - _launched_service = ServiceProcess(process, log_file) + return _service_controller + + +def start_service(): + return service_controller().start() def ensure_service(timeout=60.0): - global _service_auto_start, _service_connection, _service_snapshot - - _service_auto_start = True - deadline = time.monotonic() + timeout - last_launch = 0.0 - starting = False - while time.monotonic() < deadline: - connection, health = _discover_live_service(timeout=0.5) - if connection is not None: - _service_connection = connection - _register_service_client(connection) - _service_snapshot = _format_service_snapshot(connection, health) - _redraw_service_ui() - return connection - if not starting: - _service_snapshot = {"state": "STARTING", "message": "Starting"} - _redraw_service_ui() - starting = True - _close_finished_launched_service() - if ( - _launched_service is None - and time.monotonic() - last_launch >= 1.0 - ): - start_service() - last_launch = time.monotonic() - time.sleep(0.05) - raise RuntimeError( - f"Service did not become ready; log: {service_log_path()}" - ) - - -def submit_service_job(request): - connection = _service_connection or ensure_service() - payload = dict(request.payload) - payload["client_id"] = _client_id - _service_http_request( - connection, - "POST", - "/jobs", - payload=payload, - timeout=5.0, - ) + connection = service_controller().ensure(timeout) + _redraw_service_ui() return connection def stop_service(): - global _service_auto_start, _service_connection, _service_snapshot - global _registered_instance_id - - connection = _service_connection - if connection is None: - connection, _health = _discover_live_service(timeout=0.5) - if connection is None: - _close_launched_service(terminate=True) - _service_auto_start = False - _service_snapshot = {"state": "STOPPED", "message": "Stopped"} - return - _register_service_client(connection) - _service_http_request( - connection, - "POST", - "/shutdown", - payload={"client_id": _client_id}, - timeout=2.0, - ) - _service_auto_start = False - deadline = time.monotonic() + 10.0 - while time.monotonic() < deadline: - try: - _check_service(connection, timeout=0.2) - except RuntimeError: - break - time.sleep(0.05) - else: - raise RuntimeError("Service did not stop") - _service_connection = None - _registered_instance_id = None - _close_finished_launched_service() - _service_snapshot = {"state": "STOPPED", "message": "Stopped"} + service_controller().stop() def restart_service(): - global _service_auto_start, _service_connection, _registered_instance_id - global _service_snapshot - - connection = ensure_service() - _service_http_request( - connection, - "POST", - "/shutdown", - payload={"client_id": _client_id}, - timeout=2.0, - ) - _service_auto_start = True - _service_connection = None - _registered_instance_id = None - _service_snapshot = {"state": "STARTING", "message": "Restarting"} + return service_controller().restart() def enable_service(): - global _service_auto_start, _service_connection, _service_snapshot - - _service_auto_start = True - connection, health = _discover_live_service(timeout=0.2) - if connection is not None: - _service_connection = connection - _register_service_client(connection) - _service_snapshot = _format_service_snapshot(connection, health) - _redraw_service_ui() - return - start_service() - _service_snapshot = {"state": "STARTING", "message": "Starting"} - _redraw_service_ui() + return service_controller().enable() def clear_service_models(): - connection = ensure_service() - return _service_http_request( - connection, - "POST", - "/models/clear", - timeout=5.0, - ) - - -def cancel_service_job(job_id): - connection = ensure_service() - return _service_http_request( - connection, - "DELETE", - f"/jobs/{job_id}", - payload={"client_id": _client_id}, - timeout=2.0, - ) + return ensure_service().client.clear_resource("models") def detach_service_client(): - global _service_connection, _registered_instance_id - - connection = _service_connection - if connection is not None and _registered_instance_id is not None: - try: - _service_http_request( - connection, - "DELETE", - f"/clients/{_client_id}", - timeout=0.5, - ) - except RuntimeError: - pass - _service_connection = None - _registered_instance_id = None - # 当前 Blender 启动的 Service 随 Blender 关闭，不等待空闲超时。 - _close_launched_service(terminate=True) + service_controller().detach() def schedule_service_start(): @@ -599,104 +343,17 @@ def unschedule_service(): def service_status(): - return dict(_service_snapshot) + return dict(service_controller().snapshot) def service_busy(): - return _service_snapshot.get("state") == "BUSY" - - -def mark_service_job_complete(job_id): - """Release a stale local busy snapshot after its terminal status is observed.""" - global _service_snapshot - - active = _service_snapshot.get("active_job") - if active and active.get("job_id") != job_id: - return - if _service_snapshot.get("state") != "BUSY": - return - _service_snapshot = { - **_service_snapshot, - "busy": False, - "active_job": None, - "state": "READY", - "message": "Ready", - } + return service_controller().snapshot.get("state") == "BUSY" def _service_timer(): - global _service_connection, _service_snapshot, _registered_instance_id - - if not env_installed(): - _service_snapshot = { - "state": "UNAVAILABLE", - "message": "Environment is not installed", - } - return 2.0 - connection, health = _discover_live_service(timeout=0.2) - if connection is None: - _service_connection = None - _registered_instance_id = None - _close_finished_launched_service() - if not _service_auto_start: - if _service_snapshot.get("state") != "STOPPED": - _service_snapshot = {"state": "STOPPED", "message": "Stopped"} - _redraw_service_ui() - return 2.0 - try: - start_service() - except OSError as error: - _service_snapshot = {"state": "ERROR", "message": str(error)} - return 2.0 - _service_snapshot = {"state": "STARTING", "message": "Starting"} - _redraw_service_ui() - return 1.0 - try: - _service_connection = connection - if _registered_instance_id == connection.instance_id: - try: - _service_http_request( - connection, - "POST", - f"/clients/{_client_id}/heartbeat", - timeout=0.5, - ) - except RuntimeError: - _registered_instance_id = None - _register_service_client(connection) - health = _check_service(connection, timeout=0.5) - _service_snapshot = _format_service_snapshot(connection, health) - except RuntimeError as error: - _service_connection = None - _registered_instance_id = None - _service_snapshot = {"state": "ERROR", "message": str(error)} + interval = service_controller().poll(available=env_installed()) _redraw_service_ui() - return 5.0 - - -def _register_service_client(connection): - global _registered_instance_id - - if _registered_instance_id == connection.instance_id: - return - _service_http_request( - connection, - "POST", - "/clients", - payload={"client_id": _client_id, "label": _client_label}, - timeout=2.0, - ) - _registered_instance_id = connection.instance_id - - -def _format_service_snapshot(connection, health): - state = "BUSY" if health.get("busy") else "READY" - return { - **health, - "state": state, - "message": "Busy" if state == "BUSY" else "Ready", - "port": connection.port, - } + return interval def _redraw_service_ui(): @@ -710,155 +367,92 @@ def _redraw_service_ui(): pass -def _close_finished_launched_service(): - if _launched_service is None or _launched_service.process.poll() is None: - return - _close_launched_service() +def download_model(model_key): + get_downloadable_model(model_key) + return JobRequest("download-model", {"model": model_key}) -def _close_launched_service(terminate=False): - global _launched_service - - launched = _launched_service - _launched_service = None - if launched is None: - return - if terminate and launched.process.poll() is None: - launched.process.terminate() - try: - launched.process.wait(timeout=3.0) - except subprocess.TimeoutExpired: - launched.process.kill() - launched.process.wait(timeout=3.0) - launched.log_file.close() - - -def download_model_command(model_key, status_path): - model = get_downloadable_model(model_key) - command = [ - str(venv_python()), - str(Path(__file__).with_name("worker.py")), - "download", - "--repository", - model.repository, - "--model-dir", - str(model_directory(model_key)), - "--status", - str(status_path), - "--label", - model.label, - ] - for pattern in model.download_patterns: - command.extend(["--allow-pattern", pattern]) - for pattern in model.ready_patterns: - command.extend(["--required-pattern", pattern]) - if model.r2_directory: - command.extend(["--r2-directory", model.r2_directory]) - for filename, size, checksum in model.r2_files: - command.extend( - ["--r2-file", f"{filename}|{size}|{checksum}"] - ) - return command - - -def remove_background_request( - settings, - input_path, - status_path, - result_directory, - model_key, - video_frames=None, - start_frame=1, - device=None, -): - return JobRequest( - { - "job_id": job_id_from_status(status_path), - "command": "remove-background", - "parameters": { - "input": str(input_path), - "model_dir": str(model_directory(model_key)), - "output_dir": str(result_directory), - "status": str(status_path), - "device": (device or settings.device).lower(), - "video_frames": ( - settings.video_frames - if video_frames is None - else video_frames - ), - "start_frame": start_frame, - }, - } +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 extract_frames_request( - settings, - input_path, - status_path, - result_directory, - video_frames=None, - start_frame=1, -): - return JobRequest( - { - "job_id": job_id_from_status(status_path), - "command": "extract-frames", - "parameters": { - "input": str(input_path), - "output_dir": str(result_directory), - "status": str(status_path), - "video_frames": ( - settings.video_frames - if video_frames is None - else video_frames - ), - "start_frame": start_frame, - }, - } +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 reconstruct_request( - settings, - export_format, - status_path, - result_directory, - input_path=None, - model_key=None, - process_res=None, - process_res_method=None, - use_ray_pose=None, - remove_background=False, - mesh_filter_edges=None, - start_frame=1, - device=None, - inference_mode="joint", -): - snapform_selected_key = model_key or selected_model_key(settings, export_format) - snapform_selected_process_res = ( - settings.process_res if process_res is None else process_res - ) - snapform_selected_process_res_method = ( - settings.process_res_method - if process_res_method is None - else process_res_method - ) - snapform_selected_ray_pose = ( - settings.use_ray_pose if use_ray_pose is None else use_ray_pose - ) - parameters = { - "input": str( - Path(input_path) - if input_path is not None - else Path(bpy.path.abspath(settings.input_path)) +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, ), - "model_dir": str(model_directory(snapform_selected_key)), - "output_dir": str(result_directory), "export_format": export_format, - "status": str(status_path), - "process_res": snapform_selected_process_res, - "process_res_method": snapform_selected_process_res_method, + "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, @@ -866,83 +460,52 @@ def reconstruct_request( "video_frames": settings.video_frames, "start_frame": start_frame, "device": (device or settings.device).lower(), - "use_ray_pose": 1 if snapform_selected_ray_pose else 0, "ref_view_strategy": settings.ref_view_strategy, - "inference_mode": inference_mode, - "remove_background_model_dir": "", - "mesh_filter_edges": bool( - getattr(settings, "mesh_filter_edges", True) - if mesh_filter_edges is None - else mesh_filter_edges - ), - "mesh_max_edge_length": float( - getattr(settings, "mesh_max_edge_length", 0.1) - ), } - if remove_background: - parameters["remove_background_model_dir"] = str( - model_directory(BEN2_MODEL.key) - ) - return JobRequest( - { - "job_id": job_id_from_status(status_path), - "command": "reconstruct", - "parameters": parameters, - } - ) -def moge_request( - settings, - status_path, - result_directory, - input_path=None, - resolution_level=None, - include_points=False, - remove_background=False, - start_frame=1, - device=None, - model_key=None, - debug_export_format="", - surface_mesh=False, -): - from .models import BEN2_MODEL, MOGE2_MODEL +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 "", + }) - selected_model_key = model_key or MOGE2_MODEL.key - parameters = { - "input": str( - Path(input_path) - if input_path is not None - else Path(bpy.path.abspath(settings.input_path)) +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, ), - "model_dir": str(model_directory(selected_model_key)), - "output_dir": str(result_directory), - "status": str(status_path), - "resolution_level": int( - settings.moge_resolution_level - if resolution_level is None - else resolution_level - ), - "include_points": bool(include_points), - "debug_export_format": debug_export_format, - "surface_mesh": bool(surface_mesh), + "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(), - "remove_background_model_dir": ( - str(model_directory(BEN2_MODEL.key)) - if remove_background - else "" - ), } - return JobRequest( - { - "job_id": job_id_from_status(status_path), - "command": "moge", - "parameters": parameters, - } - ) def creation_flags(): @@ -951,7 +514,7 @@ def creation_flags(): return subprocess.CREATE_NO_WINDOW -def read_status(path): +def read_bootstrap_status(path): try: return json.loads(Path(path).read_text(encoding="utf-8")) except (FileNotFoundError, json.JSONDecodeError, OSError): diff --git a/src/snapform/install_env.py b/src/snapform/install_env.py index a9599c0..daf077c 100644 --- a/src/snapform/install_env.py +++ b/src/snapform/install_env.py @@ -391,7 +391,7 @@ def install_stages(args, venv_dir, python, uv): ), ] models = DOWNLOADABLE_MODELS - worker_script = Path(__file__).with_name("worker.py") + 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] diff --git a/src/snapform/operators.py b/src/snapform/operators.py index a7d1cf4..acdca04 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -1,13 +1,17 @@ -import os import shutil -import subprocess -import threading 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, @@ -17,36 +21,6 @@ from .models import ( ) -RECONSTRUCTION_JOB = "RECONSTRUCT" -IMAGE_DEPTH_JOB = "IMAGE_DEPTH" -MOGE2_DEPTH_JOB = "MOGE2_DEPTH" -MOGE2_MESH_JOB = "MOGE2_MESH" -MOGE2_DEBUG_JOB = "MOGE2_DEBUG" -IMAGE_MESH_JOB = "IMAGE_MESH" -BACKGROUND_REMOVAL_JOB = "REMOVE_BACKGROUND" -IMAGE_BACKGROUND_JOB = "IMAGE_BACKGROUND" -ORIGINAL_BACKGROUND_JOB = "ORIGINAL_BACKGROUND" -ORIGINAL_FRAMES_JOB = "ORIGINAL_FRAMES" -DA3_IMAGE_JOBS = frozenset({IMAGE_DEPTH_JOB, IMAGE_MESH_JOB}) -IMAGE_EMPTY_JOBS = frozenset( - { - *DA3_IMAGE_JOBS, - MOGE2_DEPTH_JOB, - MOGE2_MESH_JOB, - IMAGE_BACKGROUND_JOB, - ORIGINAL_BACKGROUND_JOB, - ORIGINAL_FRAMES_JOB, - } -) -PRODUCTION_AI_JOBS = frozenset( - { - *DA3_IMAGE_JOBS, - MOGE2_DEPTH_JOB, - MOGE2_MESH_JOB, - IMAGE_BACKGROUND_JOB, - ORIGINAL_BACKGROUND_JOB, - } -) IMAGE_DEPTH_DEFAULT_PROCESS_RES = 1024 IMAGE_DEPTH_PROCESS_RES_METHOD = "upper_bound_resize" BEN2_IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"} @@ -72,7 +46,7 @@ def is_ben2_input(input_path): def empty_image_start_frame(source_object): """Return the first frame to extract, honoring the Empty's Offset within the clip length.""" - from .image_depth import extract_frame_range + from .blender.depth_plane import extract_frame_range return extract_frame_range(source_object)[0] @@ -87,13 +61,10 @@ def open_job_directory(result_directory): @dataclass -class ActiveJob: +class SnapFormJobData: + """SnapForm-only values needed when importing a completed result.""" + kind: str - process: subprocess.Popen | None - log_file: object - log_path: Path - status_path: Path - result_directory: Path export_format: str = "" source_object_name: str = "" depth_data_name: str = "" @@ -103,436 +74,103 @@ class ActiveJob: remove_background: bool = False mesh_create_camera: bool = False temporary_input_path: Path | None = None - cancelled: bool = False - uses_service: bool = False - started: bool = False - start_error: str = "" - start_thread: threading.Thread | None = None - progress_stage: int | None = None -_active_job = None +class SnapFormJobOperator(ModalJobOperatorBase): + """SnapForm request validation and completed-result integration.""" - -def active_job(): - return _active_job - - -def cancel_active_job(): - if _active_job is None: - return - _active_job.cancelled = True - if not getattr(_active_job, "started", True): - return - if _active_job.uses_service: - try: - client.cancel_service_job( - client.job_id_from_status(_active_job.status_path) - ) - except RuntimeError: - pass - elif _active_job.process is not None and _active_job.process.poll() is None: - _active_job.process.terminate() - - -def close_active_job(): - global _active_job - - if _active_job is None: - return - if _active_job.log_file is not None: - _active_job.log_file.close() - if _active_job.temporary_input_path is not None: - temporary_path = _active_job.temporary_input_path - try: - 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 - _active_job = None - - -class RunJob(bpy.types.Operator): - bl_idname = "snapform.run_job" - bl_label = "Run SnapForm Job" bl_options = {"INTERNAL"} + result_kind = "" + export_format = "" + starting_message = "Starting task" + uses_temporary_input = False + uses_production_device = False - job_type: bpy.props.EnumProperty( - items=( - ("INSTALL_ENV", "Install Environment", ""), - ("DOWNLOAD_MODEL", "Download Model", ""), - (RECONSTRUCTION_JOB, "Export DA3 Result", ""), - (IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""), - (MOGE2_DEPTH_JOB, "Convert Image with MoGe-2", ""), - (MOGE2_MESH_JOB, "Convert Image to Surface with MoGe-2", ""), - (MOGE2_DEBUG_JOB, "Test MoGe-2 API", ""), - (IMAGE_MESH_JOB, "Convert Image to Surface", ""), - (BACKGROUND_REMOVAL_JOB, "Test BEN2 Background Removal", ""), - (IMAGE_BACKGROUND_JOB, "Remove Image Empty Background", ""), - (ORIGINAL_BACKGROUND_JOB, "Build Original Background Plane", ""), - (ORIGINAL_FRAMES_JOB, "Extract Original Depth Frames", ""), - ), - options={"HIDDEN", "SKIP_SAVE"}, - ) - model_key: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_input_path: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_source_object_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_depth_data_name: bpy.props.StringProperty( - options={"HIDDEN", "SKIP_SAVE"}, - ) - 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"}, - ) - snapform_delete_source: bpy.props.BoolProperty( - default=True, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_delete_input: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_unshaded: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_mesh_create_camera: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_remove_background: bpy.props.BoolProperty( - default=False, - options={"HIDDEN", "SKIP_SAVE"}, - ) - snapform_moge_resolution_level: bpy.props.IntProperty( - default=5, - min=0, - max=9, - options={"HIDDEN", "SKIP_SAVE"}, - ) + 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 execute(self, context): - global _active_job - - if active_job() is not None: - self.report({"WARNING"}, "Another SnapForm task is running") - return {"CANCELLED"} + 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 - try: - command, paths = self._prepare_command(context, settings) - uses_service = isinstance(command, client.JobRequest) - log_path = client.service_log_path() if uses_service else paths["log"] - working_directory = client.storage_root() - creation_flags = client.creation_flags() - except (OSError, RuntimeError, ValueError) as error: - self.report({"ERROR"}, str(error)) - return {"CANCELLED"} + if self.result_kind != "install_environment" and not client.env_installed(): + raise RuntimeError("Install the SnapForm Environment first") + return self.create_request(settings) - export_format = "" - if self.job_type == RECONSTRUCTION_JOB: - export_format = settings.export_format - elif self.job_type == IMAGE_DEPTH_JOB: - export_format = "mini_npz" - elif self.job_type == IMAGE_MESH_JOB: - export_format = "mesh_data" - elif self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: - export_format = "moge_npz" - - job = ActiveJob( - kind=self.job_type, - process=None, - log_file=None, - log_path=log_path, - status_path=paths["status"], - result_directory=paths["result"], - export_format=export_format, - source_object_name=self.snapform_source_object_name, - depth_data_name=self.snapform_depth_data_name, - process_res=self.snapform_process_res, - delete_source=self.snapform_delete_source, - unshaded=self.snapform_unshaded, + 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, - ), + mesh_create_camera=getattr(self, "snapform_mesh_create_camera", False), temporary_input_path=( Path(self.snapform_input_path) - if self.job_type in IMAGE_EMPTY_JOBS - and self.snapform_delete_input + if self.uses_temporary_input + and getattr(self, "snapform_delete_input", False) else None ), - uses_service=uses_service, ) - _active_job = job - settings.status = self._starting_message(settings) - settings.progress = 0.0 - self._timer = context.window_manager.event_timer_add( - 0.25, - window=context.window, - ) - context.window_manager.modal_handler_add(self) - self._redraw_ui(context, force=True) - job.start_thread = threading.Thread( - target=self._start_job, - args=( - job, - command, - working_directory, - creation_flags, - ), - name="SnapFormJobStart", - daemon=True, - ) - job.start_thread.start() - return {"RUNNING_MODAL"} - def _start_job(self, job, command, working_directory, creation_flags): - if job.cancelled: - job.started = True - return - try: - if getattr(job, "uses_service", False): - client.submit_service_job(command) - else: - if job.kind == "INSTALL_ENV": - client.stop_service() - job.log_file = job.log_path.open("w", encoding="utf-8") - process_environment = None - if job.kind == "INSTALL_ENV": - process_environment = os.environ.copy() - process_environment["PYTHONIOENCODING"] = ( - "utf-8:backslashreplace" - ) - process_environment["PYTHONUTF8"] = "1" - job.process = subprocess.Popen( - command, - stdout=job.log_file, - stderr=subprocess.STDOUT, - cwd=working_directory, - creationflags=creation_flags, - env=process_environment, - ) - job.started = True - if job.cancelled: - if job.uses_service: - client.cancel_service_job( - client.job_id_from_status(job.status_path) - ) - elif job.process is not None: - job.process.terminate() - except (OSError, RuntimeError) as error: - if job.log_file is not None: - job.log_file.close() - job.log_file = None - job.start_error = str(error) + def create_request(self, _settings): + raise NotImplementedError - def _prepare_command(self, _context, settings): - paths = client.new_job_paths(self.job_type) - if ( - getattr(self, "snapform_delete_input", False) - and self.job_type in IMAGE_EMPTY_JOBS - ): - archived_input = client.move_input_into_job( - self.snapform_input_path, - paths["result"], - ) - self.snapform_input_path = str(archived_input) - self.snapform_delete_input = False - remove_background = bool( - getattr(self, "snapform_remove_background", False) - ) - if self.job_type == "INSTALL_ENV": - self._require_online_access() - return client.install_env_command(paths["status"]), paths - - if not client.env_installed(): - raise RuntimeError("Install the SnapForm Environment first") - - if self.job_type == "DOWNLOAD_MODEL": - self._require_online_access() - model = get_downloadable_model(self.model_key) - return ( - client.download_model_command(model.key, paths["status"]), - paths, - ) - - if self.job_type in IMAGE_EMPTY_JOBS: - configured_input = self.snapform_input_path - elif self.job_type == BACKGROUND_REMOVAL_JOB: - configured_input = settings.background_input_path - else: - configured_input = settings.input_path + def _input_path(self, configured_input): input_path = Path(bpy.path.abspath(configured_input)) - if not (input_path.is_file() or input_path.is_dir()): - raise ValueError("Select a valid image, video, or frame directory") - start_frame = 1 - if getattr(self, "snapform_source_object_name", ""): - source_object = bpy.data.objects.get( - self.snapform_source_object_name - ) - if source_object is not None: - start_frame = empty_image_start_frame(source_object) - if self.job_type in { - BACKGROUND_REMOVAL_JOB, - IMAGE_BACKGROUND_JOB, - ORIGINAL_BACKGROUND_JOB, - }: - if not is_ben2_input(input_path): - raise ValueError( - "BEN2 requires an image, video, or directory of image frames" - ) - if not client.model_ready(BEN2_MODEL.key): - raise RuntimeError(f"Download {BEN2_MODEL.label} first") - return ( - client.remove_background_request( - settings, - input_path, - paths["status"], - paths["result"], - BEN2_MODEL.key, - start_frame=start_frame, - device=self._task_device(settings), - ), - paths, - ) - if self.job_type == ORIGINAL_FRAMES_JOB: - return ( - client.extract_frames_request( - settings, - input_path, - paths["status"], - paths["result"], - start_frame=start_frame, - ), - paths, - ) - if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: - model_key = ( - settings.moge_model - if self.job_type == MOGE2_DEBUG_JOB - else self.snapform_model_key or MOGE2_MODEL.key - ) - model = get_downloadable_model(model_key) - if not client.model_ready(model_key): - raise RuntimeError(f"Download {model.label} first") - if remove_background and not client.model_ready(BEN2_MODEL.key): - raise RuntimeError(f"Download {BEN2_MODEL.label} first") - return ( - client.moge_request( - settings, - paths["status"], - paths["result"], - input_path=input_path, - resolution_level=( - self.snapform_moge_resolution_level - if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB} - else settings.moge_resolution_level - ), - include_points=self.job_type in {MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}, - remove_background=( - self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB} - and remove_background - ), - start_frame=start_frame, - device=self._task_device(settings), - model_key=model_key, - debug_export_format=( - settings.moge_export_format - if self.job_type == MOGE2_DEBUG_JOB - else "" - ), - surface_mesh=self.job_type == MOGE2_MESH_JOB, - ), - paths, - ) - if self.job_type not in { - RECONSTRUCTION_JOB, - IMAGE_DEPTH_JOB, - IMAGE_MESH_JOB, - }: - raise RuntimeError(f"Unsupported task type: {self.job_type}") - export_format = settings.export_format - if self.job_type == IMAGE_DEPTH_JOB: - export_format = "mini_npz" - elif self.job_type == IMAGE_MESH_JOB: - export_format = "mesh_data" - use_ray_pose = None - process_res = None - model_key = selected_model_key(settings, export_format) - if self.job_type in DA3_IMAGE_JOBS: - model_key = self.snapform_model_key - use_ray_pose = self.snapform_use_ray_pose - process_res = self.snapform_process_res - model = get_model(model_key) + 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") - if ( - self.job_type in DA3_IMAGE_JOBS - and remove_background - and not client.model_ready(BEN2_MODEL.key) - ): - raise RuntimeError(f"Download {BEN2_MODEL.label} first") - if export_format == "gs_ply" and not model.supports_gs: - raise RuntimeError( - "PLY export requires a DA3 Giant or Nested Giant model" + 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" ) - return ( - client.reconstruct_request( - settings, - export_format, - paths["status"], - paths["result"], - input_path=input_path, - model_key=model_key, - process_res=process_res, - process_res_method=( - IMAGE_DEPTH_PROCESS_RES_METHOD - if self.job_type in DA3_IMAGE_JOBS - else None - ), - use_ray_pose=use_ray_pose, - start_frame=start_frame, - remove_background=( - self.job_type in DA3_IMAGE_JOBS - and remove_background - ), - device=self._task_device(settings), - inference_mode=( - "per_frame" - if self.job_type in DA3_IMAGE_JOBS - else "joint" - ), - ), - paths, - ) + self._require_ben2() + return input_path def _task_device(self, settings): """Debug panels use the scene device; production tasks use Preferences.""" - if self.job_type in PRODUCTION_AI_JOBS: + if self.uses_production_device: return client.production_device() return settings.device @@ -543,271 +181,78 @@ class RunJob(bpy.types.Operator): ) def _starting_message(self, settings): - labels = { - "INSTALL_ENV": "Starting Environment installation", - "DOWNLOAD_MODEL": "Starting model download", - } - if self.job_type == RECONSTRUCTION_JOB: - return f"Starting {settings.export_format} export" - if self.job_type == IMAGE_DEPTH_JOB: - return "Starting depth estimation" - if self.job_type in {MOGE2_DEPTH_JOB, MOGE2_MESH_JOB, MOGE2_DEBUG_JOB}: - return "Starting MoGe-2 inference" - if self.job_type == IMAGE_MESH_JOB: - return "Starting perspective surface generation" - if self.job_type == BACKGROUND_REMOVAL_JOB: - return "Starting BEN2 background removal" - if self.job_type == IMAGE_BACKGROUND_JOB: - return "Starting Image Empty background removal" - if self.job_type == ORIGINAL_BACKGROUND_JOB: - return "Starting background removal" - if self.job_type == ORIGINAL_FRAMES_JOB: - return "Extracting source frames" - if self.job_type == "DOWNLOAD_MODEL": - model = get_downloadable_model(self.model_key) - return f"Starting {model.label} download" - return labels[self.job_type] + return self.starting_message - def modal(self, context, event): - if event.type == "ESC": - cancel_active_job() + def _export_format(self, _settings): + return self.export_format - if event.type != "TIMER": - return {"PASS_THROUGH"} + def update_job_ui(self, context, progress, message): + settings = context.scene.snapform_settings + settings.progress = progress + settings.status = message - job = active_job() - if job is None: - self._remove_timer(context) - return {"CANCELLED"} + def redraw_job_ui(self, context, force=False): + from .status_bar import redraw_ui - self._update_status(context, job) - if job.start_error: - settings = context.scene.snapform_settings - settings.status = "Task failed to start" - settings.progress = 0.0 - self.report( - {"ERROR"}, - f"Unable to start task: {job.start_error}", - ) - self._remove_timer(context) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - if not job.started: - return {"RUNNING_MODAL"} - if job.uses_service: - status = client.read_status(job.status_path) or {} - state = status.get("state") - if state == "succeeded": - return self._finish(context, job, 0) - if state in {"failed", "cancelled"}: - return self._finish(context, job, 1) - if client.service_status().get("state") in {"ERROR", "STOPPED"}: - return self._finish(context, job, 1) - return {"RUNNING_MODAL"} - return_code = job.process.poll() - if return_code is None: - return {"RUNNING_MODAL"} - return self._finish(context, job, return_code) + redraw_ui(context, force=force) - def _update_status(self, context, job): - status = client.read_status(job.status_path) - if status is None: + 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 - settings = context.scene.snapform_settings - reported_progress = float( - status.get("progress", settings.progress) - ) - reported_progress = min(max(reported_progress, 0.0), 1.0) - reported_stage = status.get("stage") - if reported_stage is not None and reported_stage != getattr( - job, - "progress_stage", - None, - ): - job.progress_stage = reported_stage - settings.progress = reported_progress - else: - settings.progress = max(settings.progress, reported_progress) - settings.status = self._format_status(status, settings.status) - self._redraw_ui(context, force=True) + 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 _format_status(self, status, fallback): - message = status.get("message") - stage = status.get("stage") - stages = status.get("stages") - stage_label = status.get("stage_label") - if stage and stages: - prefix = f"{stage}/{stages}" - normalized = str(message or "").strip() - if normalized and normalized.lower() != str( - stage_label or "" - ).lower(): - return f"{prefix} {normalized}" - if stage_label: - return f"{prefix} {stage_label}" - return prefix - if message: - return str(message).strip() - return str(stage_label or fallback) - - def _finish(self, context, job, return_code): - settings = context.scene.snapform_settings - self._remove_timer(context) - if job.log_file is not None: - job.log_file.close() - - if job.cancelled: - settings.status = "Task cancelled" - settings.progress = 0.0 - if getattr(job, "uses_service", False): - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - - status = client.read_status(job.status_path) or {} - error = status.get("error") - if return_code != 0 or error: - settings.status = "Task failed" - settings.progress = 0.0 - fallback_error = ( - "Service task failed" - if job.uses_service - else "Task subprocess failed" - ) - self.report( - {"ERROR"}, - f"{error or fallback_error}; log: {job.log_path}", - ) - if job.uses_service: - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - - if job.kind in { - RECONSTRUCTION_JOB, - BACKGROUND_REMOVAL_JOB, - MOGE2_DEBUG_JOB, - }: + 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 == RECONSTRUCTION_JOB: + if job.kind == "debug_da3": message = ( f"Exported {job.export_format} to {job.result_directory}" ) - elif job.kind == BACKGROUND_REMOVAL_JOB: + 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) - if job.kind == IMAGE_DEPTH_JOB: - try: - settings.status = "Building depth plane" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._build_depth_plane(context, job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "Depth import failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "Depth plane is ready" - if job.kind == MOGE2_DEPTH_JOB: - try: - settings.status = "Building MoGe-2 depth plane" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._build_moge2_depth_plane(job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "MoGe-2 depth import failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete( - client.job_id_from_status(job.status_path) - ) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "MoGe-2 depth plane is ready" - if job.kind == MOGE2_MESH_JOB: - try: - settings.status = "Building MoGe-2 surface" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._build_moge2_image_mesh(job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "MoGe-2 surface import failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "MoGe-2 surface is ready" - if job.kind == IMAGE_MESH_JOB: - try: - settings.status = "Building perspective surface" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._build_image_mesh(job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "Surface import failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "Perspective surface is ready" - if job.kind in {ORIGINAL_BACKGROUND_JOB, ORIGINAL_FRAMES_JOB}: - try: - settings.status = "Building depth plane" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._build_original_depth_plane(job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "Depth import failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "Depth plane is ready" - if job.kind == IMAGE_BACKGROUND_JOB: - try: - settings.status = "Updating Image Empty" - settings.progress = 0.96 - self._redraw_ui(context, force=True) - self._apply_removed_background(job) - except (OSError, RuntimeError, ValueError) as import_error: - settings.status = "Image update failed" - settings.progress = 0.0 - self.report({"ERROR"}, str(import_error)) - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"CANCELLED"} - status["message"] = "Image background is removed" - - settings.status = status.get("message", "Task complete") - settings.progress = 1.0 - try: - job.status_path.unlink(missing_ok=True) - except OSError: - pass - if getattr(job, "uses_service", False): - client.mark_service_job_complete(client.job_id_from_status(job.status_path)) - close_active_job() - self._redraw_ui(context, force=True) - return {"FINISHED"} + 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( @@ -881,7 +326,7 @@ class RunJob(bpy.types.Operator): is_sequence = bool(list(result_directory.glob("frame.*"))) background_path = "" color_frames_dir = "" - if job.kind == ORIGINAL_BACKGROUND_JOB: + if job.kind == "generate_original_background": background_path = str( result_directory if is_sequence @@ -954,22 +399,308 @@ class RunJob(bpy.types.Operator): if "FINISHED" not in undo_result: raise RuntimeError("Unable to create the background removal Undo step") - def _remove_timer(self, context): - timer = getattr(self, "_timer", None) - if timer is None: - return - context.window_manager.event_timer_remove(timer) - self._timer = None +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 _redraw_ui(self, context, force=False): - from .status_bar import redraw_ui + def _image_input(self): + return self._input_path(self.snapform_input_path) - redraw_ui(context, force=force) - def cancel(self, context): - cancel_active_job() - self._remove_timer(context) - self._redraw_ui(context, force=True) +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) + ) + 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), + ) + + +class Moge2ImageJobOperator(GeometryImageJobOperator): + request_name = "" + snapform_model_key: bpy.props.StringProperty( + default="BASE", + options={"HIDDEN", "SKIP_SAVE"}, + ) + snapform_moge_resolution_level: bpy.props.IntProperty( + default=5, + min=0, + max=9, + options={"HIDDEN", "SKIP_SAVE"}, + ) + + 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) + ) + 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), + ) + + +class BackgroundRequestMixin: + request_name = "" + + 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), + ) + + +class InstallEnvironment(SnapFormJobOperator): + bl_idname = "snapform.install_environment" + bl_label = "Install Environment" + result_kind = "install_environment" + starting_message = "Starting Environment installation" + + def create_request(self, _settings): + self._require_online_access() + return JobRequest("install-environment") + + +class DownloadModel(SnapFormJobOperator): + bl_idname = "snapform.download_model" + bl_label = "Download Model" + result_kind = "download_model" + model_key: bpy.props.StringProperty( + options={"HIDDEN", "SKIP_SAVE"}, + ) + + def create_request(self, _settings): + self._require_online_access() + return client.download_model(self.model_key) + + def _starting_message(self, _settings): + model = get_downloadable_model(self.model_key) + return f"Starting {model.label} download" + + +class DebugDa3(SnapFormJobOperator): + bl_idname = "snapform.debug_da3" + bl_label = "Debug DA3" + result_kind = "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 _starting_message(self, settings): + return f"Starting {settings.export_format} export" + + def _export_format(self, settings): + return settings.export_format + + +class DebugMoge2(SnapFormJobOperator): + bl_idname = "snapform.debug_moge2" + bl_label = "Debug MoGe-2" + result_kind = "debug_moge2" + export_format = "moge_npz" + 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), + ) + + +class DebugBen2(BackgroundRequestMixin, SnapFormJobOperator): + bl_idname = "snapform.debug_ben2" + bl_label = "Debug BEN2" + result_kind = "debug_ben2" + request_name = "debug_ben2" + starting_message = "Starting BEN2 background removal" + + +class GenerateDa3Depth(Da3ImageJobOperator): + 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" + starting_message = "Starting depth estimation" + + +class GenerateMoge2Depth(Moge2ImageJobOperator): + 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" + starting_message = "Starting MoGe-2 inference" + + +class GenerateDa3Surface(Da3ImageJobOperator): + 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" + starting_message = "Starting perspective surface generation" + + +class GenerateMoge2Surface(Moge2ImageJobOperator): + 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" + starting_message = "Starting MoGe-2 inference" + + +class ProcessImageBackground(BackgroundRequestMixin, ImageInputJobOperator): + bl_idname = "snapform.process_image_background" + bl_label = "Process Image Background" + result_kind = "process_image_background" + request_name = "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 + + +class GenerateOriginalBackground( + BackgroundRequestMixin, + GeometryImageJobOperator, +): + bl_idname = "snapform.generate_original_background" + bl_label = "Generate Original Background" + result_kind = "generate_original_background" + request_name = "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 + + +class ExtractOriginalFrames(GeometryImageJobOperator): + bl_idname = "snapform.extract_original_frames" + bl_label = "Extract Original Frames" + result_kind = "extract_original_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(), + ) + + +JOB_OPERATOR_CLASSES = ( + InstallEnvironment, + DownloadModel, + DebugDa3, + DebugMoge2, + DebugBen2, + GenerateDa3Depth, + GenerateMoge2Depth, + GenerateDa3Surface, + GenerateMoge2Surface, + ProcessImageBackground, + GenerateOriginalBackground, + ExtractOriginalFrames, +) class CancelJob(bpy.types.Operator): diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 76d8fd4..0856814 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -24,12 +24,11 @@ class AiServicePanel(bpy.types.Panel): ) install_row = layout.row() install_row.enabled = not running - operator = install_row.operator( - "snapform.run_job", + install_row.operator( + "snapform.install_environment", text="Install Environment", icon="IMPORT", ) - operator.job_type = "INSTALL_ENV" return status = client.service_status() @@ -46,7 +45,7 @@ class AiServicePanel(bpy.types.Panel): if state in {"READY", "BUSY"}: row.label( text=( - f"Running (Clients: {status.get('clients', 0)})" + "Running" if state == "READY" else "Running (Busy)" ), @@ -54,7 +53,6 @@ class AiServicePanel(bpy.types.Panel): ) if state == "READY": controls = row.row(align=True) - controls.enabled = status.get("clients", 0) <= 1 operator = controls.operator( "snapform.control_service", text="", @@ -88,7 +86,11 @@ class AiServicePanel(bpy.types.Panel): operator.action = "START" if state in {"READY", "BUSY"}: - cached = status.get("cached_models", {}) + cached = ( + status.get("resources", {}) + .get("models", {}) + .get("loaded", {}) + ) cached_names = [name for name in cached.values() if name] cached_row = layout.row(align=True) cached_row.label( @@ -142,11 +144,10 @@ class ModelsPanel(bpy.types.Panel): download_row = row.row(align=True) download_row.enabled = env_installed() and not running operator = download_row.operator( - "snapform.run_job", + "snapform.download_model", text="Download", icon="IMPORT", ) - operator.job_type = "DOWNLOAD_MODEL" operator.model_key = model.key @@ -231,12 +232,11 @@ class Da3OutputDebugPanel(bpy.types.Panel): and not client.service_busy() and model_ready(selected_model_key(settings)) ) - operator = export_row.operator( - "snapform.run_job", + export_row.operator( + "snapform.debug_da3", text="Generate Output", icon="EXPORT", ) - operator.job_type = "RECONSTRUCT" class Moge2ApiDebugPanel(bpy.types.Panel): @@ -262,12 +262,11 @@ class Moge2ApiDebugPanel(bpy.types.Panel): and not client.service_busy() and model_ready(settings.moge_model) ) - operator = run_row.operator( - "snapform.run_job", + run_row.operator( + "snapform.debug_moge2", text="Generate MoGe-2 Output", icon="EXPORT", ) - operator.job_type = "MOGE2_DEBUG" class Ben2RemovalDebugPanel(bpy.types.Panel): @@ -291,9 +290,8 @@ class Ben2RemovalDebugPanel(bpy.types.Panel): and not client.service_busy() and model_ready(BEN2_MODEL.key) ) - operator = background_row.operator( - "snapform.run_job", + background_row.operator( + "snapform.debug_ben2", text="Remove Background", icon="IMAGE_ALPHA", ) - operator.job_type = "REMOVE_BACKGROUND" diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index 88948cd..0ee79ec 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -41,12 +41,11 @@ class SnapFormPreferences(bpy.types.AddonPreferences): else: install_row = layout.row() install_row.enabled = active_job() is None - operator = install_row.operator( - "snapform.run_job", + install_row.operator( + "snapform.install_environment", text="Install Environment", icon="IMPORT", ) - operator.job_type = "INSTALL_ENV" layout.separator() operator = layout.operator( "snapform.control_service", diff --git a/src/snapform/service.py b/src/snapform/service.py deleted file mode 100644 index 537f793..0000000 --- a/src/snapform/service.py +++ /dev/null @@ -1,598 +0,0 @@ -import argparse -import gc -import hmac -import json -import os -import sys -import threading -import time -import traceback -from contextlib import asynccontextmanager -from pathlib import Path -from types import SimpleNamespace - - -SCRIPT_DIRECTORY = Path(__file__).resolve().parent -if str(SCRIPT_DIRECTORY) not in sys.path: - sys.path.insert(0, str(SCRIPT_DIRECTORY)) - -import worker - - -CLIENT_TIMEOUT_SECONDS = 10.0 -JOB_COMMANDS = frozenset( - {"remove-background", "reconstruct", "extract-frames", "moge"} -) -IDLE_SHUTDOWN_SECONDS = 20.0 - - -class ServiceLock: - def __init__(self, path): - self.path = Path(path) - self.file = None - - def acquire(self): - self.path.parent.mkdir(parents=True, exist_ok=True) - lock_file = self.path.open("a+b") - lock_file.seek(0, os.SEEK_END) - if lock_file.tell() == 0: - lock_file.write(b"0") - lock_file.flush() - lock_file.seek(0) - try: - if sys.platform == "win32": - import msvcrt - - msvcrt.locking(lock_file.fileno(), msvcrt.LK_NBLCK, 1) - else: - import fcntl - - fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB) - except OSError: - lock_file.close() - return False - self.file = lock_file - return True - - def release(self): - if self.file is None: - return - try: - self.file.seek(0) - if sys.platform == "win32": - import msvcrt - - msvcrt.locking(self.file.fileno(), msvcrt.LK_UNLCK, 1) - else: - import fcntl - - fcntl.flock(self.file.fileno(), fcntl.LOCK_UN) - finally: - self.file.close() - self.file = None - - -class ModelCache: - def __init__(self): - self.ben2_model = None - self.ben2_key = None - self.da3_model = None - self.da3_key = None - self.moge_model = None - self.moge_key = None - - def get_ben2(self, model_dir, device): - requested_device = getattr(device, "type", device) - key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.ben2_model is not None and self.ben2_key == key: - return self.ben2_model, 0.0 - self.release_ben2() - return self._load_ben2(model_dir, requested_device, key) - - def _load_ben2(self, model_dir, requested_device, key): - from ben2_onnx import create_session - - started = time.perf_counter() - model = create_session(model_dir, requested_device) - self.ben2_model = model - self.ben2_key = key - elapsed_ms = (time.perf_counter() - started) * 1000.0 - return model, elapsed_ms - - def get_da3(self, model_dir, device): - requested_device = getattr(device, "type", device) - key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.moge_model is not None: - self.release_moge() - if self.da3_model is not None and self.da3_key == key: - return self.da3_model - self.release_da3() - return self._load_da3(model_dir, requested_device, key) - - def _load_da3(self, model_dir, requested_device, key): - from da3_onnx 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 - - key = (str(Path(model_dir).resolve()), requested_device, "onnx") - if self.da3_model is not None: - self.release_da3() - if self.moge_model is not None and self.moge_key == key: - return self.moge_model - self.release_moge() - model = create_session(model_dir, requested_device) - self.moge_model = model - self.moge_key = key - return model - - def release_ben2(self): - self.ben2_model = None - self.ben2_key = None - self._release_cuda_cache() - - def release_da3(self): - self.da3_model = None - self.da3_key = None - self._release_cuda_cache() - - def release_moge(self): - self.moge_model = None - self.moge_key = None - self._release_cuda_cache() - - def close(self): - self.ben2_model = None - self.da3_model = None - self.moge_model = None - self.ben2_key = None - self.da3_key = None - self.moge_key = None - self._release_cuda_cache() - - def status(self): - return { - "ben2": ( - Path(self.ben2_key[0]).name if self.ben2_key is not None else "" - ), - "da3": ( - Path(self.da3_key[0]).name - if self.da3_key is not None - else "" - ), - "moge2": ( - Path(getattr(self, "moge_key", ("",))[0]).name - if getattr(self, "moge_key", None) is not None - else "" - ), - } - - def _release_cuda_cache(self): - gc.collect() - - -class ServiceState: - def __init__(self, token, instance_id): - self.token = token - self.instance_id = instance_id - self.models = ModelCache() - self.lock = threading.Lock() - self.clients = {} - self.active_job = None - self.started_at = time.time() - self.last_without_clients = time.monotonic() - self.shutdown_event = threading.Event() - self.monitor_stop = threading.Event() - self.monitor_thread = None - - def register_client(self, client_id, label): - with self.lock: - self.clients[client_id] = { - "label": label, - "last_seen": time.monotonic(), - } - self.last_without_clients = None - - def heartbeat(self, client_id): - with self.lock: - client = self.clients.get(client_id) - if client is None: - return False - client["last_seen"] = time.monotonic() - return True - - def unregister_client(self, client_id): - with self.lock: - self.clients.pop(client_id, None) - self._mark_empty_clients() - - def reserve(self, job_id, client_id, command, status_path): - with self.lock: - if client_id not in self.clients or self.active_job is not None: - return False - self.active_job = { - "job_id": job_id, - "client_id": client_id, - "command": command, - "status": status_path, - "cancel_requested": False, - } - return True - - def snapshot(self): - self.expire_clients() - with self.lock: - return { - "ready": True, - "instance_id": self.instance_id, - "busy": self.active_job is not None, - "active_job": dict(self.active_job) if self.active_job else None, - "clients": len(self.clients), - "cached_models": self.models.status(), - "started_at": self.started_at, - } - - def clear_models(self): - with self.lock: - if self.active_job is not None: - return False - self.models.close() - return True - - def request_shutdown(self, client_id): - with self.lock: - if client_id not in self.clients: - return "missing" - if len(self.clients) > 1: - return "shared" - if self.active_job is not None: - return "busy" - self.shutdown_event.set() - return "accepted" - - def cancel_job(self, job_id, client_id): - with self.lock: - active = self.active_job - if active is None or active["job_id"] != job_id: - return None - if active["client_id"] != client_id: - return None - active["cancel_requested"] = True - return dict(active) - - def is_job_cancelled(self, job_id): - with self.lock: - active = self.active_job - return bool( - active is not None - and active["job_id"] == job_id - and active["cancel_requested"] - ) - - def run_job(self, payload): - job_id = payload["job_id"] - parameters = payload["parameters"] - status_path = parameters["status"] - worker.write_status( - status_path, - 0.0, - "Service accepted the task", - state="running", - ) - try: - cancel_check = lambda: self.is_job_cancelled(job_id) - worker.check_cancelled(cancel_check) - args = SimpleNamespace(**parameters) - if payload["command"] == "remove-background": - worker.remove_background( - args, - model_cache=self.models, - cancel_check=cancel_check, - ) - elif payload["command"] == "extract-frames": - worker.extract_frames( - args, - cancel_check=cancel_check, - ) - elif payload["command"] == "moge": - worker.run_moge( - args, - model_cache=self.models, - cancel_check=cancel_check, - ) - else: - worker.reconstruct( - args, - model_cache=self.models, - cancel_check=cancel_check, - ) - worker.check_cancelled(cancel_check) - except worker.TaskCancelled: - worker.write_status( - status_path, - 1.0, - "Task cancelled", - state="cancelled", - ) - except Exception as error: - traceback.print_exc() - worker.write_status( - status_path, - 1.0, - "Task failed", - str(error), - state="failed", - ) - finally: - with self.lock: - if ( - self.active_job is not None - and self.active_job["job_id"] == job_id - ): - self.active_job = None - self._mark_empty_clients() - - def start_monitor(self): - self.monitor_thread = threading.Thread( - target=self._monitor, - name="SnapFormServiceMonitor", - daemon=True, - ) - self.monitor_thread.start() - - def close(self): - self.monitor_stop.set() - if self.monitor_thread is not None: - self.monitor_thread.join(timeout=1.0) - self.models.close() - - def expire_clients(self): - deadline = time.monotonic() - CLIENT_TIMEOUT_SECONDS - with self.lock: - expired = [ - client_id - for client_id, client in self.clients.items() - if client["last_seen"] < deadline - ] - for client_id in expired: - self.clients.pop(client_id, None) - self._mark_empty_clients() - - def _mark_empty_clients(self): - if self.clients: - self.last_without_clients = None - elif self.last_without_clients is None: - self.last_without_clients = time.monotonic() - - def _monitor(self): - while not self.monitor_stop.wait(1.0): - self.expire_clients() - with self.lock: - idle = self.active_job is None - empty_since = self.last_without_clients - if not idle or empty_since is None: - continue - if time.monotonic() - empty_since >= IDLE_SHUTDOWN_SECONDS: - self.shutdown_event.set() - return - - -def create_app(token, instance_id="test"): - from fastapi import BackgroundTasks, FastAPI, Header, HTTPException - - state = ServiceState(token, instance_id) - - @asynccontextmanager - async def lifespan(_app): - state.start_monitor() - yield - state.close() - - app = FastAPI( - title="SnapForm Service", - docs_url=None, - redoc_url=None, - openapi_url=None, - lifespan=lifespan, - ) - app.state.snapform_service_state = state - - def authorize(x_snapform_token: str = Header(default="")): - if not hmac.compare_digest(x_snapform_token, state.token): - raise HTTPException(status_code=401, detail="Invalid service token") - - @app.get("/health") - def health(x_snapform_token: str = Header(default="")): - authorize(x_snapform_token) - return state.snapshot() - - @app.post("/clients", status_code=201) - def register_client( - payload: dict, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - client_id = str(payload.get("client_id", "")) - if not client_id: - raise HTTPException(status_code=422, detail="Missing client id") - state.register_client(client_id, str(payload.get("label", "Blender"))) - return {"client_id": client_id, "lease_seconds": CLIENT_TIMEOUT_SECONDS} - - @app.post("/clients/{client_id}/heartbeat") - def heartbeat( - client_id: str, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - if not state.heartbeat(client_id): - raise HTTPException(status_code=404, detail="Client is not registered") - return {"client_id": client_id, "state": "active"} - - @app.delete("/clients/{client_id}") - def unregister_client( - client_id: str, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - state.unregister_client(client_id) - return {"client_id": client_id, "state": "removed"} - - @app.post("/jobs", status_code=202) - def submit_job( - payload: dict, - background_tasks: BackgroundTasks, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - job_id = str(payload.get("job_id", "")) - client_id = str(payload.get("client_id", "")) - command = payload.get("command") - parameters = payload.get("parameters") - if not job_id or command not in JOB_COMMANDS: - raise HTTPException(status_code=422, detail="Invalid job request") - if not isinstance(parameters, dict) or not parameters.get("status"): - raise HTTPException(status_code=422, detail="Missing job parameters") - if not state.reserve( - job_id, - client_id, - command, - parameters["status"], - ): - raise HTTPException(status_code=409, detail="Service is busy") - worker.write_status( - parameters["status"], - 0.0, - "Task queued", - state="queued", - ) - background_tasks.add_task(state.run_job, payload) - return {"job_id": job_id, "state": "queued"} - - @app.delete("/jobs/{job_id}", status_code=202) - def cancel_job( - job_id: str, - payload: dict, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - client_id = str(payload.get("client_id", "")) - active = state.cancel_job(job_id, client_id) - if active is None: - raise HTTPException(status_code=404, detail="Active job was not found") - try: - current = json.loads( - Path(active["status"]).read_text(encoding="utf-8") - ) - progress = float(current.get("progress", 0.0)) - except (OSError, ValueError, TypeError, json.JSONDecodeError): - progress = 0.0 - worker.write_status( - active["status"], - progress, - "Cancelling task", - state="cancelling", - ) - return {"job_id": job_id, "state": "cancelling"} - - @app.post("/models/clear") - def clear_models(x_snapform_token: str = Header(default="")): - authorize(x_snapform_token) - if not state.clear_models(): - raise HTTPException(status_code=409, detail="Service is busy") - return {"state": "cleared"} - - @app.post("/shutdown", status_code=202) - def shutdown_service( - payload: dict, - x_snapform_token: str = Header(default=""), - ): - authorize(x_snapform_token) - result = state.request_shutdown(str(payload.get("client_id", ""))) - if result == "shared": - raise HTTPException(status_code=409, detail="Other Blender clients are connected") - if result == "busy": - raise HTTPException(status_code=409, detail="Service is busy") - if result != "accepted": - raise HTTPException(status_code=404, detail="Client is not registered") - return {"state": "stopping"} - - return app - - -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("--token", required=True) - parser.add_argument("--instance-id", required=True) - parser.add_argument("--discovery", required=True) - parser.add_argument("--lock", required=True) - return parser.parse_args() - - -def write_discovery(args): - path = Path(args.discovery) - path.parent.mkdir(parents=True, exist_ok=True) - temporary = path.with_name(f"{path.name}.tmp") - temporary.write_text( - json.dumps( - { - "port": args.port, - "token": args.token, - "pid": os.getpid(), - "instance_id": args.instance_id, - } - ), - encoding="utf-8", - ) - temporary.replace(path) - - -def remove_discovery(path, instance_id): - path = Path(path) - try: - payload = json.loads(path.read_text(encoding="utf-8")) - except (FileNotFoundError, json.JSONDecodeError, OSError): - return - if payload.get("instance_id") != instance_id: - return - path.unlink(missing_ok=True) - - -def main(): - import uvicorn - - args = parse_args() - service_lock = ServiceLock(args.lock) - if not service_lock.acquire(): - return - try: - app = create_app(args.token, args.instance_id) - config = uvicorn.Config( - app, - host=args.host, - port=args.port, - access_log=False, - log_level="warning", - ) - server = uvicorn.Server(config) - - def watch_shutdown(): - app.state.snapform_service_state.shutdown_event.wait() - server.should_exit = True - - threading.Thread(target=watch_shutdown, daemon=True).start() - write_discovery(args) - server.run() - finally: - remove_discovery(args.discovery, args.instance_id) - service_lock.release() - - -if __name__ == "__main__": - main() diff --git a/tests/test_addon.py b/tests/test_addon.py index dc8aac0..9a17289 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -86,7 +86,7 @@ class AddonTest(unittest.TestCase): self.status_bar_draws.clear() self.context_menu_draws.clear() self.outliner_menu_draws.clear() - self.addon.operators._active_job = None + self.addon.operators.close_active_job() self.fake_bpy.data = SimpleNamespace( objects=SimpleNamespace(get=lambda _name: None), ) @@ -153,6 +153,34 @@ class AddonTest(unittest.TestCase): 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 = [] @@ -204,9 +232,8 @@ class AddonTest(unittest.TestCase): self.assertEqual(len(calls), 1) identifier, options, operator = calls[0] - self.assertEqual(identifier, "snapform.run_job") + self.assertEqual(identifier, "snapform.install_environment") self.assertEqual(options["text"], "Install Environment") - self.assertEqual(operator.job_type, "INSTALL_ENV") def test_image_mesh_reference_plane_matches_image_empty_bounds(self): source_object = SimpleNamespace( @@ -510,7 +537,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + process_image_background=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -527,7 +554,7 @@ class AddonTest(unittest.TestCase): result = operator.execute(context) self.assertEqual(result, {"FINISHED"}) - self.assertEqual(calls[0]["job_type"], "IMAGE_BACKGROUND") + self.assertNotIn("job_type", calls[0]) self.assertEqual(calls[0]["snapform_input_path"], str(input_path)) self.assertEqual( calls[0]["snapform_source_object_name"], @@ -549,7 +576,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + generate_original_background=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -565,7 +592,7 @@ class AddonTest(unittest.TestCase): result = operator.execute(context) self.assertEqual(result, {"FINISHED"}) - self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND") + self.assertNotIn("job_type", calls[0]) self.assertNotIn("snapform_background_depth_distance", calls[0]) def test_original_depth_animation_extracts_frames_and_background(self): @@ -589,7 +616,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + generate_original_background=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -606,7 +633,7 @@ class AddonTest(unittest.TestCase): self.assertEqual(result, {"FINISHED"}) self.assertEqual(scene.snapform_settings.video_frames, 97) - self.assertEqual(calls[0]["job_type"], "ORIGINAL_BACKGROUND") + 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"]) @@ -626,7 +653,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + extract_original_frames=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -642,7 +669,7 @@ class AddonTest(unittest.TestCase): result = operator.execute(context) self.assertEqual(result, {"FINISHED"}) - self.assertEqual(calls[0]["job_type"], "ORIGINAL_FRAMES") + self.assertNotIn("job_type", calls[0]) self.assertTrue(calls[0]["snapform_delete_source"]) def test_conversion_dialog_exposes_requested_settings(self): @@ -744,7 +771,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + process_image_background=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -788,7 +815,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + generate_da3_depth=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -805,7 +832,7 @@ class AddonTest(unittest.TestCase): self.assertEqual(result, {"FINISHED"}) self.assertEqual(scene.snapform_settings.video_frames, 24) - self.assertEqual(calls[0]["job_type"], "IMAGE_DEPTH") + self.assertNotIn("job_type", calls[0]) self.assertTrue(calls[0]["snapform_remove_background"]) self.assertNotIn("snapform_background_depth_distance", calls[0]) @@ -836,10 +863,10 @@ class AddonTest(unittest.TestCase): self.assertEqual(labels["MOGE2_VITB_NORMAL"], "MoGe-2 ViT-B Normal (Missing)") def test_update_status_formats_stage_progress(self): - operator = self.addon.RunJob() + operator = self.addon.operators.ModalJobOperatorBase() self.assertEqual( - operator._format_status( + operator.format_job_status( { "progress": 0.5, "message": "2 downloaded", @@ -852,7 +879,7 @@ class AddonTest(unittest.TestCase): "3/8 2 downloaded", ) self.assertEqual( - operator._format_status( + operator.format_job_status( { "progress": 0.0, "message": "", @@ -865,7 +892,7 @@ class AddonTest(unittest.TestCase): "6/8 da3 small", ) self.assertEqual( - operator._format_status( + operator.format_job_status( { "progress": 0.5, "message": "Estimating depth", @@ -875,7 +902,7 @@ class AddonTest(unittest.TestCase): "Estimating depth", ) self.assertEqual( - operator._format_status( + operator.format_job_status( { "progress": 0.25, "message": "Downloading packages", @@ -889,7 +916,7 @@ class AddonTest(unittest.TestCase): "2/3 Downloading packages", ) self.assertEqual( - operator._format_status( + operator.format_job_status( { "progress": 0.5, "message": "", @@ -902,26 +929,43 @@ class AddonTest(unittest.TestCase): "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) + job_type = "example" + + 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.RunJob() + operator = self.addon.operators.SnapFormJobOperator() settings = SimpleNamespace(progress=0.9, status="previous") context = SimpleNamespace( scene=SimpleNamespace(snapform_settings=settings) ) - job = SimpleNamespace(status_path=Path("status.json"), progress_stage=1) + job = SimpleNamespace(progress_stage=1, progress=0.9) - with patch.object( - self.addon.client, - "read_status", - return_value={ + 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", - }, - ), patch.object(operator, "_redraw_ui"): - operator._update_status(context, job) + }, + ) self.assertEqual(job.progress_stage, 2) self.assertEqual(settings.progress, 0.1) @@ -962,7 +1006,7 @@ class AddonTest(unittest.TestCase): calls = [] self.fake_bpy.ops = SimpleNamespace( snapform=SimpleNamespace( - run_job=lambda **options: calls.append(options) + generate_da3_depth=lambda **options: calls.append(options) or {"RUNNING_MODAL"}, ), ) @@ -1021,38 +1065,32 @@ class AddonTest(unittest.TestCase): self.assertEqual(image_settings.file_format, "OPEN_EXR") def test_temporary_packed_input_is_removed_when_job_closes(self): - closed = [] with TemporaryDirectory() as directory: input_path = Path(directory) / "packed.png" input_path.touch() - self.addon.operators._active_job = SimpleNamespace( - log_file=SimpleNamespace(close=lambda: closed.append(True)), + data = self.addon.operators.SnapFormJobData( + kind="test", temporary_input_path=input_path, ) - self.addon.operators.close_active_job() + self.addon.operators.SnapFormJobOperator().cleanup_job_data(data) self.assertFalse(input_path.exists()) - self.assertEqual(closed, [True]) - self.assertIsNone(self.addon.operators.active_job()) def test_temporary_frame_directory_is_removed_when_job_closes(self): - closed = [] with TemporaryDirectory() as directory: root = Path(directory) temporary_directory = root / "inputs" / "frames" temporary_directory.mkdir(parents=True) (temporary_directory / "frame-0000.jpg").touch() - self.addon.operators._active_job = SimpleNamespace( - log_file=SimpleNamespace(close=lambda: closed.append(True)), + data = self.addon.operators.SnapFormJobData( + kind="test", temporary_input_path=temporary_directory, ) - self.addon.operators.close_active_job() + self.addon.operators.SnapFormJobOperator().cleanup_job_data(data) self.assertFalse(temporary_directory.exists()) - self.assertEqual(closed, [True]) - self.assertIsNone(self.addon.operators.active_job()) def test_prepare_image_input_passes_movie_file_through(self): with TemporaryDirectory() as directory: @@ -1108,8 +1146,7 @@ class AddonTest(unittest.TestCase): model="BASE", gs_model="GIANT_1_1", ) - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator = self.addon.operators.GenerateDa3Depth() operator.snapform_source_object_name = "Reference" operator.snapform_process_res = 336 @@ -1129,10 +1166,6 @@ class AddonTest(unittest.TestCase): ("NESTED_GIANT_LARGE_1_1", True), ): with self.subTest(model=snapform_model), patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -1142,41 +1175,26 @@ class AddonTest(unittest.TestCase): return_value=True, ) as model_ready, patch.object( self.addon.client, - "reconstruct_request", - return_value=self.addon.client.JobRequest( + "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, prepared_paths = operator._prepare_command( - None, - settings, - ) + command = operator.create_request(settings) self.assertIsInstance( command, - self.addon.client.JobRequest, + self.addon.operators.JobRequest, ) - self.assertEqual(prepared_paths, paths) model_ready.assert_called_once_with(snapform_model) - self.assertEqual(reconstruct.call_args.args[1], "mini_npz") - self.assertEqual( - reconstruct.call_args.kwargs["input_path"], - input_path, - ) - self.assertEqual( - reconstruct.call_args.kwargs["model_key"], - 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["process_res_method"], - "upper_bound_resize", - ) self.assertEqual( reconstruct.call_args.kwargs["use_ray_pose"], snapform_ray_pose, @@ -1189,8 +1207,7 @@ class AddonTest(unittest.TestCase): model="BASE", gs_model="GIANT_1_1", ) - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator = self.addon.operators.GenerateDa3Depth() operator.snapform_model_key = "BASE" operator.snapform_use_ray_pose = True operator.snapform_process_res = 1024 @@ -1201,14 +1218,6 @@ class AddonTest(unittest.TestCase): input_path.touch() operator.snapform_input_path = str(input_path) with patch.object( - self.addon.client, - "new_job_paths", - return_value={ - "log": Path(directory) / "job.log", - "status": Path(directory) / "job.json", - "result": Path(directory) / "result", - }, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -1218,7 +1227,7 @@ class AddonTest(unittest.TestCase): side_effect=lambda key: key == "BASE", ): with self.assertRaisesRegex(RuntimeError, "BEN2 Base"): - operator._prepare_command(None, settings) + operator.create_request(settings) def test_image_depth_background_option_reaches_service_request(self): client = importlib.import_module("snapform.client") @@ -1246,17 +1255,18 @@ class AddonTest(unittest.TestCase): "model_directory", side_effect=lambda key: Path(key.lower()), ): - request = client.reconstruct_request( + request = importlib.import_module("snapform.client").generate_da3_depth( settings, - "mini_npz", - Path("status.json"), - Path("result"), + Path("input.png"), + "BASE", + process_res=504, + use_ray_pose=False, remove_background=True, ) self.assertEqual( - request.payload["parameters"]["remove_background_model_dir"], - str(Path("ben2_base")), + request.payload["parameters"]["remove_background_model"], + "BEN2_BASE", ) self.assertNotIn( "background_depth_distance", @@ -1265,23 +1275,23 @@ class AddonTest(unittest.TestCase): def test_service_busy_snapshot_is_released_for_completed_local_job(self): client = self.addon.client - previous_snapshot = client._service_snapshot - client._service_snapshot = { + controller = client.service_controller() + previous_snapshot = controller.snapshot + controller.snapshot = { "state": "BUSY", "busy": True, "active_job": {"job_id": "depth-job"}, } try: - client.mark_service_job_complete("depth-job") + controller.mark_job_complete("depth-job") self.assertEqual(client.service_status()["state"], "READY") self.assertFalse(client.service_busy()) finally: - client._service_snapshot = previous_snapshot + controller.snapshot = previous_snapshot def test_job_submission_starts_without_blocking_execute(self): - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator = self.addon.operators.GenerateDa3Depth() operator.snapform_source_object_name = "Reference" operator.snapform_depth_data_name = "img.depth" operator.snapform_process_res = 1024 @@ -1306,20 +1316,15 @@ class AddonTest(unittest.TestCase): window_manager=window_manager, window=object(), ) - paths = { - "log": Path("job.log"), - "status": Path("job.json"), - "result": Path("result"), - } - request = self.addon.client.JobRequest({"job_id": "job"}) + request = self.addon.operators.JobRequest({"job_id": "job"}) try: with patch.object( operator, - "_prepare_command", - return_value=(request, paths), + "create_job_request", + return_value=request, ), patch.object( - self.addon.client, - "submit_service_job", + self.addon.client.service_controller(), + "submit", side_effect=submit, ), patch.object( self.addon.client, @@ -1333,7 +1338,7 @@ class AddonTest(unittest.TestCase): self.addon.client, "creation_flags", return_value=0, - ), patch.object(operator, "_redraw_ui"): + ), patch.object(operator, "redraw_job_ui"): before = time.perf_counter() result = operator.execute(context) elapsed = time.perf_counter() - before @@ -1372,28 +1377,10 @@ class AddonTest(unittest.TestCase): def test_closing_blender_terminates_service_it_started(self): client = importlib.import_module("snapform.client") - previous_connection = client._service_connection - previous_registered = client._registered_instance_id - client._service_connection = SimpleNamespace(port=12345) - client._registered_instance_id = "instance" - try: - with patch.object( - client, - "_service_http_request", - ) as request, patch.object( - client, - "_close_launched_service", - ) as close_service: - client.detach_service_client() - - request.assert_called_once() - self.assertEqual(request.call_args.args[1], "DELETE") - close_service.assert_called_once_with(terminate=True) - finally: - client._service_connection = previous_connection - client._registered_instance_id = previous_registered - self.assertIsNone(client._service_connection) - self.assertIsNone(client._registered_instance_id) + 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__[ @@ -1412,6 +1399,7 @@ class AddonTest(unittest.TestCase): 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, @@ -1434,15 +1422,14 @@ class AddonTest(unittest.TestCase): "model_directory", return_value=Path("model"), ): - request = client.reconstruct_request( + request = importlib.import_module("snapform.client").debug_da3( settings, - "gs_ply", - Path("status.json"), - Path("result"), + Path("input.png"), + "NESTED_GIANT_LARGE_1_1", ) parameters = request.payload["parameters"] - self.assertEqual(request.payload["command"], "reconstruct") + self.assertEqual(request.payload["job_type"], "debug-da3") self.assertEqual(parameters["inference_mode"], "joint") self.assertEqual(parameters["export_format"], "gs_ply") for parameter in ( @@ -1470,11 +1457,9 @@ class AddonTest(unittest.TestCase): "model_directory", return_value=Path("ben2_base"), ): - request = client.remove_background_request( + request = importlib.import_module("snapform.client").remove_background( settings, Path("clip.mp4"), - Path("status.json"), - Path("result"), "BEN2_BASE", start_frame=70, ) @@ -1488,101 +1473,19 @@ class AddonTest(unittest.TestCase): client = importlib.import_module("snapform.client") settings = SimpleNamespace(device="AUTO", video_frames=24) - request = client.extract_frames_request( + request = importlib.import_module("snapform.client").extract_frames( settings, Path("clip.mp4"), - Path("status.json"), - Path("result"), start_frame=70, ) parameters = request.payload["parameters"] - self.assertEqual(request.payload["command"], "extract-frames") + self.assertEqual(request.payload["job_type"], "extract-frames") self.assertEqual(parameters["input"], str(Path("clip.mp4"))) - self.assertEqual(parameters["output_dir"], str(Path("result"))) + self.assertNotIn("output_dir", parameters) self.assertEqual(parameters["video_frames"], 24) self.assertEqual(parameters["start_frame"], 70) - def test_new_job_paths_wraps_each_job_in_a_folder(self): - client = importlib.import_module("snapform.client") - with TemporaryDirectory() as directory: - jobs_root = Path(directory) - with patch.object( - client, - "jobs_directory", - return_value=jobs_root, - ): - paths = client.new_job_paths("RECONSTRUCT") - - self.assertTrue(paths["result"].is_dir()) - self.assertEqual(paths["result"].parent, jobs_root) - self.assertEqual(paths["log"].parent, paths["result"]) - self.assertEqual(paths["status"].parent, paths["result"]) - self.assertEqual(paths["log"].name, "job.log") - self.assertEqual(paths["status"].name, "status.json") - self.assertEqual( - client.job_id_from_status(paths["status"]), - paths["result"].name, - ) - - def test_new_job_paths_uses_unique_action_object_names(self): - client = importlib.import_module("snapform.client") - with TemporaryDirectory() as directory: - jobs_root = Path(directory) - with patch.object( - client, - "jobs_directory", - return_value=jobs_root, - ): - depth_paths = client.new_job_paths( - self.addon.operators.IMAGE_DEPTH_JOB - ) - background_paths = client.new_job_paths( - self.addon.operators.IMAGE_BACKGROUND_JOB - ) - frames_paths = client.new_job_paths( - self.addon.operators.ORIGINAL_FRAMES_JOB - ) - - self.assertRegex( - depth_paths["result"].name, - r"^generate-depth-plane_\d{8}-\d{6}-\d{6}$", - ) - self.assertTrue( - background_paths["result"].name.startswith( - "remove-image-background_" - ) - ) - self.assertTrue( - frames_paths["result"].name.startswith( - "extract-original-frames_" - ) - ) - - def test_generated_inputs_are_moved_into_job_root(self): - client = importlib.import_module("snapform.client") - with TemporaryDirectory() as directory: - root = Path(directory) - staging = root / "staging" - staging.mkdir() - (staging / "frame.0001.jpg").write_bytes(b"first") - (staging / "frame.0002.jpg").write_bytes(b"second") - job_directory = root / "job" - job_directory.mkdir() - - result = client.move_input_into_job(staging, job_directory) - - self.assertEqual(result, job_directory) - self.assertFalse(staging.exists()) - self.assertEqual( - (job_directory / "frame.0001.jpg").read_bytes(), - b"first", - ) - self.assertEqual( - (job_directory / "frame.0002.jpg").read_bytes(), - b"second", - ) - def test_empty_image_start_frame_honors_image_user(self): operators = self.addon.operators @@ -1636,8 +1539,7 @@ class AddonTest(unittest.TestCase): use_ray_pose=False, ref_view_strategy="saddle_balanced", ) - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB + operator = self.addon.operators.GenerateDa3Depth() operator.snapform_source_object_name = "Reference" operator.snapform_model_key = "BASE" operator.snapform_use_ray_pose = True @@ -1662,10 +1564,6 @@ class AddonTest(unittest.TestCase): "result": Path(directory) / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -1675,20 +1573,16 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.client, - "reconstruct_request", - return_value=self.addon.client.JobRequest( + "generate_da3_depth", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as reconstruct: - operator._prepare_command(None, settings) + operator.create_request(settings) finally: self.fake_bpy.data.objects.get = original_get self.assertEqual(reconstruct.call_args.kwargs["start_frame"], 71) - self.assertEqual( - reconstruct.call_args.kwargs["inference_mode"], - "per_frame", - ) def test_reconstruct_request_accepts_image_depth_overrides(self): client = importlib.import_module("snapform.client") @@ -1716,23 +1610,21 @@ class AddonTest(unittest.TestCase): "model_directory", return_value=Path("large-model"), ) as snapform_model_directory: - request = client.reconstruct_request( + request = importlib.import_module("snapform.client").generate_da3_depth( settings, - "mini_npz", - Path("status.json"), - Path("result"), - model_key="BASE", + Path("input.png"), + "BASE", process_res=252, - process_res_method="lower_bound_resize", use_ray_pose=False, ) - snapform_model_directory.assert_called_once_with("BASE") + 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"], - "lower_bound_resize", + "upper_bound_resize", ) self.assertEqual(parameters["use_ray_pose"], 0) @@ -1762,11 +1654,12 @@ class AddonTest(unittest.TestCase): "model_directory", side_effect=lambda key: Path(key.lower()), ): - request = client.reconstruct_request( + request = importlib.import_module("snapform.client").generate_da3_surface( settings, - "mesh_data", - Path("status.json"), - Path("result"), + Path("input.png"), + "BASE", + process_res=504, + use_ray_pose=False, remove_background=True, ) @@ -1777,8 +1670,8 @@ class AddonTest(unittest.TestCase): self.assertNotIn("mesh_depth_smoothing", parameters) self.assertNotIn("mesh_postprocess_method", parameters) self.assertEqual( - parameters["remove_background_model_dir"], - str(Path("ben2_base")), + parameters["remove_background_model"], + "BEN2_BASE", ) def test_export_format_property_includes_all_da3_formats(self): @@ -1821,8 +1714,7 @@ class AddonTest(unittest.TestCase): use_ray_pose=False, ref_view_strategy="saddle_balanced", ) - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.RECONSTRUCTION_JOB + operator = self.addon.operators.DebugDa3() with TemporaryDirectory() as directory: input_path = Path(directory) / "input.png" @@ -1834,10 +1726,6 @@ class AddonTest(unittest.TestCase): "result": Path(directory) / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -1847,15 +1735,15 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.client, - "reconstruct_request", - return_value=self.addon.client.JobRequest( + "debug_da3", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as reconstruct: - operator._prepare_command(None, settings) + operator.create_request(settings) - self.assertEqual(reconstruct.call_args.args[1], "mesh") - self.assertEqual(reconstruct.call_args.kwargs["model_key"], "BASE") + 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__[ @@ -2006,77 +1894,31 @@ class AddonTest(unittest.TestCase): def test_ai_operation_restarts_a_manually_stopped_service(self): client = self.addon.client - connection = client.Connection(51234, "token", 42, "host") - health = { - "instance_id": "host", - "busy": False, - "clients": 1, - "cached_models": {}, - } - previous_auto_start = client._service_auto_start - previous_connection = client._service_connection - previous_snapshot = client._service_snapshot - client._service_auto_start = False - client._service_connection = None - client._service_snapshot = {"state": "STOPPED", "message": "Stopped"} - try: - with patch.object( - client, - "_discover_live_service", - return_value=(connection, health), - ), patch.object( - client, - "_register_service_client", - ), patch.object( - client, - "_check_service", - return_value=health, - ), patch.object(client, "_redraw_service_ui"): - result = client.ensure_service(timeout=0.1) - - self.assertIs(result, connection) - self.assertTrue(client._service_auto_start) - self.assertEqual(client.service_status()["state"], "READY") - finally: - client._service_auto_start = previous_auto_start - client._service_connection = previous_connection - client._service_snapshot = previous_snapshot + 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 - previous_auto_start = client._service_auto_start - previous_snapshot = client._service_snapshot - client._service_auto_start = False - client._service_snapshot = {"state": "STOPPED", "message": "Stopped"} - try: - with patch.object( - client, - "env_installed", - return_value=True, - ), patch.object( - client, - "_discover_live_service", - return_value=(None, None), - ), patch.object( - client, - "_close_finished_launched_service", - ), patch.object(client, "start_service") as start_service: - interval = client._service_timer() - - start_service.assert_not_called() - self.assertEqual(interval, 2.0) - self.assertEqual(client.service_status()["state"], "STOPPED") - finally: - client._service_auto_start = previous_auto_start - client._service_snapshot = previous_snapshot + 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.RunJob() - operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB + operator = self.addon.operators.DebugBen2() with TemporaryDirectory() as directory: root = Path(directory) @@ -2089,10 +1931,6 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -2102,29 +1940,24 @@ class AddonTest(unittest.TestCase): return_value=True, ) as model_ready, patch.object( self.addon.client, - "remove_background_request", - return_value=self.addon.client.JobRequest( + "debug_ben2", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as remove_background: - command, prepared_paths = operator._prepare_command( - None, - settings, - ) + command = operator.create_request(settings) - self.assertIsInstance(command, self.addon.client.JobRequest) - self.assertEqual(prepared_paths, paths) + 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[4], "BEN2_BASE") + 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.RunJob() - operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB + operator = self.addon.operators.DebugBen2() with TemporaryDirectory() as directory: root = Path(directory) @@ -2137,10 +1970,6 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -2150,12 +1979,12 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.client, - "remove_background_request", - return_value=self.addon.client.JobRequest( + "debug_ben2", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as remove_background: - operator._prepare_command(None, settings) + operator.create_request(settings) self.assertEqual( remove_background.call_args.kwargs["device"], @@ -2164,8 +1993,7 @@ class AddonTest(unittest.TestCase): def test_production_jobs_use_the_preferences_device(self): settings = SimpleNamespace(device="CPU", video_frames=16) - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB + operator = self.addon.operators.ProcessImageBackground() with TemporaryDirectory() as directory: root = Path(directory) @@ -2178,10 +2006,6 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -2195,12 +2019,12 @@ class AddonTest(unittest.TestCase): return_value="MPS", ), patch.object( self.addon.client, - "remove_background_request", - return_value=self.addon.client.JobRequest( + "remove_background", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as remove_background: - operator._prepare_command(None, settings) + operator.create_request(settings) self.assertEqual( remove_background.call_args.kwargs["device"], @@ -2209,8 +2033,7 @@ class AddonTest(unittest.TestCase): def test_image_background_job_uses_operator_input(self): settings = SimpleNamespace(device="AUTO") - operator = self.addon.RunJob() - operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB + operator = self.addon.operators.ProcessImageBackground() with TemporaryDirectory() as directory: root = Path(directory) @@ -2223,10 +2046,6 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -2236,24 +2055,19 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.client, - "remove_background_request", - return_value=self.addon.client.JobRequest( + "remove_background", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as remove_background: - command, prepared_paths = operator._prepare_command( - None, - settings, - ) + command = operator.create_request(settings) - self.assertIsInstance(command, self.addon.client.JobRequest) - self.assertEqual(prepared_paths, paths) + 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.RunJob() - operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB + operator = self.addon.operators.ProcessImageBackground() with TemporaryDirectory() as directory: root = Path(directory) @@ -2266,10 +2080,6 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, @@ -2279,23 +2089,19 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.client, - "remove_background_request", - return_value=self.addon.client.JobRequest( + "remove_background", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as remove_background: - command, prepared_paths = operator._prepare_command( - None, - settings, - ) + command = operator.create_request(settings) - self.assertIsInstance(command, self.addon.client.JobRequest) + 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.RunJob() - operator.job_type = self.addon.operators.ORIGINAL_FRAMES_JOB + operator = self.addon.operators.ExtractOriginalFrames() with TemporaryDirectory() as directory: root = Path(directory) @@ -2308,27 +2114,19 @@ class AddonTest(unittest.TestCase): "result": root / "result", } with patch.object( - self.addon.client, - "new_job_paths", - return_value=paths, - ), patch.object( self.addon.client, "env_installed", return_value=True, ), patch.object( self.addon.client, - "extract_frames_request", - return_value=self.addon.client.JobRequest( + "extract_frames", + return_value=self.addon.operators.JobRequest( {"job_id": "test"} ), ) as extract_frames: - command, prepared_paths = operator._prepare_command( - None, - settings, - ) + command = operator.create_request(settings) - self.assertIsInstance(command, self.addon.client.JobRequest) - self.assertEqual(prepared_paths, paths) + 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"], @@ -2349,7 +2147,7 @@ class AddonTest(unittest.TestCase): or {"FINISHED"}, ), ) - operator = self.addon.RunJob() + operator = self.addon.operators.SnapFormJobOperator() job = SimpleNamespace( result_directory=Path("result"), source_object_name="Reference", @@ -2387,7 +2185,7 @@ class AddonTest(unittest.TestCase): or {"FINISHED"}, ), ) - operator = self.addon.RunJob() + operator = self.addon.operators.SnapFormJobOperator() with TemporaryDirectory() as directory: result_directory = Path(directory) (result_directory / "foreground.0000.png").touch() @@ -2422,9 +2220,9 @@ class AddonTest(unittest.TestCase): or {"FINISHED"}, ), ) - operator = self.addon.RunJob() + operator = self.addon.operators.SnapFormJobOperator() job = SimpleNamespace( - kind=self.addon.operators.ORIGINAL_BACKGROUND_JOB, + kind="generate_original_background", result_directory=Path("result"), source_object_name="Reference", subdivision_level=5, @@ -2464,13 +2262,13 @@ class AddonTest(unittest.TestCase): or {"FINISHED"}, ), ) - operator = self.addon.RunJob() + 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=self.addon.operators.ORIGINAL_FRAMES_JOB, + kind="extract_original_frames", result_directory=result_directory, source_object_name="Reference", subdivision_level=5, @@ -2497,30 +2295,14 @@ class AddonTest(unittest.TestCase): ("undo", {"message": "Convert to Depth Plane"}), ) - def test_ben2_download_command_filters_large_alternative_weights(self): + def test_ben2_download_uses_model_job(self): client = importlib.import_module("snapform.client") - with patch.object( - client, - "venv_python", - return_value=Path("python"), - ), patch.object( - client, - "model_directory", - return_value=Path("model"), - ): - command = client.download_model_command( - "BEN2_BASE", - Path("status.json"), - ) + request = importlib.import_module("snapform.client").download_model( + "BEN2_BASE" + ) - patterns = [ - command[index + 1] - for index, value in enumerate(command) - if value == "--allow-pattern" - ] - self.assertEqual(patterns, ["onnx/model_fp16.onnx"]) - label_index = command.index("--label") - self.assertEqual(command[label_index + 1], "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") @@ -2590,34 +2372,6 @@ class AddonTest(unittest.TestCase): self.assertEqual(len(calls), 1) self.assertIn("filepath", calls[0]) - def test_completed_job_removes_status_file(self): - operator = self.addon.RunJob() - with TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - status_path.write_text( - '{"state": "succeeded", "progress": 1.0, "message": "done"}', - encoding="utf-8", - ) - job = SimpleNamespace( - kind="INSTALL_ENV", - cancelled=False, - log_file=None, - status_path=status_path, - ) - context = SimpleNamespace( - scene=SimpleNamespace( - snapform_settings=SimpleNamespace( - status="", - progress=0.0, - ), - ), - ) - with patch.object(operator, "_redraw_ui"): - result = operator._finish(context, job, 0) - - self.assertEqual(result, {"FINISHED"}) - self.assertFalse(status_path.exists()) - def test_status_bar_draws_progress_and_cancel_button(self): progress_calls = [] operator_calls = [] @@ -2637,7 +2391,6 @@ class AddonTest(unittest.TestCase): def operator(self, operator, **options): operator_calls.append((operator, options)) - self.addon.operators._active_job = SimpleNamespace() context = SimpleNamespace( scene=SimpleNamespace( snapform_settings=SimpleNamespace( @@ -2647,10 +2400,12 @@ class AddonTest(unittest.TestCase): ), ) - self.addon.draw_status_bar( - SimpleNamespace(layout=Layout()), - context, - ) + 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, diff --git a/tests/test_ben2_onnx.py b/tests/test_ben2_onnx.py index 3649104..6decedf 100644 --- a/tests/test_ben2_onnx.py +++ b/tests/test_ben2_onnx.py @@ -7,7 +7,12 @@ import numpy as np from PIL import Image -SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform" +SOURCE_DIR = ( + Path(__file__).resolve().parents[1] + / "src" + / "snapform" + / "ai_service" +) sys.path.insert(0, str(SOURCE_DIR)) import ben2_onnx diff --git a/tests/test_da3_onnx.py b/tests/test_da3_onnx.py index 4d431c3..a668274 100644 --- a/tests/test_da3_onnx.py +++ b/tests/test_da3_onnx.py @@ -8,7 +8,12 @@ import numpy as np from PIL import Image -SOURCE_DIR = Path(__file__).resolve().parents[1] / "src" / "snapform" +SOURCE_DIR = ( + Path(__file__).resolve().parents[1] + / "src" + / "snapform" + / "ai_service" +) sys.path.insert(0, str(SOURCE_DIR)) import da3_onnx diff --git a/tests/test_download_progress.py b/tests/test_download_progress.py index 6b71f01..d640c56 100644 --- a/tests/test_download_progress.py +++ b/tests/test_download_progress.py @@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_worker(): specification = importlib.util.spec_from_file_location( "snapform_download_worker", - PROJECT_ROOT / "src" / "snapform" / "worker.py", + PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) diff --git a/tests/test_image_depth.py b/tests/test_image_depth.py index 9f14f4e..ba004e6 100644 --- a/tests/test_image_depth.py +++ b/tests/test_image_depth.py @@ -26,7 +26,11 @@ def load_image_depth(): with patch.dict(sys.modules, {"bpy": fake_bpy}): specification = importlib.util.spec_from_file_location( "snapform_image_depth_test", - PROJECT_ROOT / "src" / "snapform" / "image_depth.py", + PROJECT_ROOT + / "src" + / "snapform" + / "blender" + / "depth_plane.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) diff --git a/tests/test_moge_onnx.py b/tests/test_moge_onnx.py index 1c51364..aa1f529 100644 --- a/tests/test_moge_onnx.py +++ b/tests/test_moge_onnx.py @@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_module(): specification = importlib.util.spec_from_file_location( "snapform_moge_onnx", - PROJECT_ROOT / "src" / "snapform" / "moge_onnx.py", + PROJECT_ROOT / "src" / "snapform" / "ai_service" / "moge_onnx.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 2cd388d..1d57756 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -52,13 +52,14 @@ class PackTest(unittest.TestCase): self.assertIn("__init__.py", names) self.assertIn("blender_manifest.toml", names) - self.assertIn("image_depth.py", names) + self.assertIn("blender/depth_plane.py", names) + self.assertIn("blender/surface.py", names) self.assertIn( "assets/O_SnapForm.blend", names, ) - self.assertIn("worker.py", names) - self.assertIn("service.py", names) + self.assertIn("ai_service/jobs.py", names) + self.assertIn("ai_service/server.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_service.py index 05cf6c5..c0fe2e4 100644 --- a/tests/test_service.py +++ b/tests/test_service.py @@ -13,7 +13,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_service(): specification = importlib.util.spec_from_file_location( "snapform_service", - PROJECT_ROOT / "src" / "snapform" / "service.py", + PROJECT_ROOT / "src" / "snapform" / "ai_service" / "server.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) @@ -38,8 +38,15 @@ class ServiceTest(unittest.TestCase): def setUpClass(cls): cls.service = load_service() + def setUp(self): + self.jobs_temp = tempfile.TemporaryDirectory() + self.addCleanup(self.jobs_temp.cleanup) + self.jobs_directory = Path(self.jobs_temp.name) / "jobs" + def create_cache(self): - cache = self.service.ModelCache.__new__(self.service.ModelCache) + cache = self.service.ModelSessionCache.__new__( + self.service.ModelSessionCache + ) cache.ben2_model = None cache.ben2_key = None cache.da3_model = None @@ -63,7 +70,7 @@ class ServiceTest(unittest.TestCase): self.assertIs(first, second) self.assertEqual(loads, [("model", "cuda")]) - self.assertEqual(cache.status()["ben2"], "model") + self.assertEqual(cache.snapshot()["ben2"], "model") self.assertEqual(cached_load_ms, 0.0) def test_da3_model_is_reused_for_matching_runtime_key(self): @@ -86,7 +93,7 @@ class ServiceTest(unittest.TestCase): self.assertIs(first, second) self.assertEqual(loads, [("model", "cuda")]) - self.assertEqual(cache.status()["da3"], "model") + self.assertEqual(cache.snapshot()["da3"], "model") def test_depth_model_cache_keeps_only_da3_or_moge(self): cache = self.create_cache() @@ -118,141 +125,133 @@ class ServiceTest(unittest.TestCase): self.assertIsNone(cache.da3_model) self.assertIsNotNone(cache.ben2_model) - def test_service_tracks_clients_and_rejects_shared_shutdown(self): + def test_dedicated_service_runs_only_one_job_at_a_time(self): fake_models = SimpleNamespace( - status=lambda: {"ben2": "", "da3": ""}, + snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, + clear=lambda: None, close=lambda: None, ) - with patch.object(self.service, "ModelCache", return_value=fake_models): - state = self.service.ServiceState("token", "instance") + with patch.object(self.service, "ModelManager", return_value=fake_models): + server = self.service.create_server(jobs_directory=self.jobs_directory) - state.register_client("first", "Blender A") - state.register_client("second", "Blender B") + parameters = {"model": "BASE"} + first = server.submit("debug-da3", parameters, job_id="job") + second = server.submit("debug-da3", parameters, job_id="other") - self.assertEqual(state.snapshot()["clients"], 2) - self.assertEqual(state.request_shutdown("first"), "shared") - self.assertTrue( - state.reserve("job", "first", "reconstruct", "status.json") + self.assertIsNotNone(first) + self.assertIsNone(second) + self.assertTrue(server.snapshot("instance")["busy"]) + self.assertNotIn("clients", server.snapshot("instance")) + + def test_service_registers_tasks_instead_of_model_names(self): + server = self.service.create_server(jobs_directory=self.jobs_directory) + self.assertEqual( + set(server.handlers), + { + "download-model", + "remove-background", + "extract-frames", + "generate-da3-depth", + "generate-moge2-depth", + "generate-da3-surface", + "generate-moge2-surface", + "debug-da3", + "debug-moge2", + "debug-ben2", + }, ) - self.assertFalse( - state.reserve("other", "second", "reconstruct", "other.json") - ) - - def test_service_accepts_moge_command(self): - self.assertIn("moge", self.service.JOB_COMMANDS) def test_cancelling_job_keeps_service_and_model_cache_alive(self): fake_models = SimpleNamespace( - status=lambda: {"ben2": "BEN2_BASE", "da3": ""}, + snapshot=lambda: { + "loaded": {"ben2": "BEN2_BASE", "da3": ""} + }, + clear=lambda: None, close=lambda: None, ) - with patch.object(self.service, "ModelCache", return_value=fake_models): - state = self.service.ServiceState("token", "instance") - state.register_client("client", "Blender") - with tempfile.TemporaryDirectory() as directory: - status_path = Path(directory) / "status.json" - state.reserve( - "job", - "client", - "remove-background", - str(status_path), - ) - - active = state.cancel_job("job", "client") - self.assertTrue(active["cancel_requested"]) - self.assertTrue(state.is_job_cancelled("job")) - state.run_job( - { - "job_id": "job", - "command": "remove-background", - "parameters": {"status": str(status_path)}, - } - ) - status = self.service.json.loads( - status_path.read_text(encoding="utf-8") - ) + with patch.object(self.service, "ModelManager", return_value=fake_models): + server = self.service.create_server(jobs_directory=self.jobs_directory) + context = server.submit("remove-background", {}, job_id="job") + self.assertTrue(context.request_cancel()) + self.assertTrue(context.is_cancelled()) + server.run(context) + status = context.snapshot() self.assertEqual(status["state"], "cancelled") - self.assertFalse(state.shutdown_event.is_set()) - self.assertEqual(state.snapshot()["cached_models"]["ben2"], "BEN2_BASE") - - def test_expired_clients_are_removed(self): - fake_models = SimpleNamespace( - status=lambda: {"ben2": "", "da3": ""}, - close=lambda: None, - ) - with patch.object(self.service, "ModelCache", return_value=fake_models): - state = self.service.ServiceState("token", "instance") - state.register_client("client", "Blender") - state.clients["client"]["last_seen"] -= ( - self.service.CLIENT_TIMEOUT_SECONDS + 1.0 + self.assertFalse(server.shutdown_event.is_set()) + self.assertEqual( + server.snapshot("instance")["resources"]["models"]["loaded"]["ben2"], + "BEN2_BASE", ) - state.expire_clients() + def test_job_progress_is_kept_in_service_memory(self): + server = self.service.JobServer("Test", jobs_directory=self.jobs_directory) - self.assertEqual(state.snapshot()["clients"], 0) + @server.job("example") + def example(context, _parameters): + context.update(0.5, "Half way") - def test_discovery_is_only_removed_by_matching_service(self): - with tempfile.TemporaryDirectory() as directory: - path = Path(directory) / "service.json" - args = SimpleNamespace( - discovery=str(path), - port=51234, - token="token", - instance_id="owner", - ) - with patch.object(self.service.os, "getpid", return_value=42): - self.service.write_discovery(args) + context = server.submit("example", {}) + server.run(context) - payload = self.service.json.loads(path.read_text(encoding="utf-8")) - self.assertEqual(payload["port"], 51234) - self.assertEqual(payload["pid"], 42) - self.service.remove_discovery(path, "other") - self.assertTrue(path.is_file()) - self.service.remove_discovery(path, "owner") - self.assertFalse(path.exists()) + 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_job_server_owns_and_clears_shared_resources(self): + events = [] + resource = SimpleNamespace( + snapshot=lambda: {"loaded": "model"}, + 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.job("example") + def example(context, _parameters): + self.assertIs(context.resource("models"), resource) + + context = server.submit("example", {}, job_id="job") + self.assertFalse(server.clear_resource("models")) + server.run(context) + self.assertTrue(server.clear_resource("models")) + server.close() + + self.assertEqual( + server.snapshot("instance")["resources"]["models"]["loaded"], + "model", + ) + self.assertEqual(events, ["clear", "close"]) def test_service_dispatches_frame_extraction_without_models(self): fake_models = SimpleNamespace( - status=lambda: {"ben2": "", "da3": ""}, + snapshot=lambda: {"loaded": {"ben2": "", "da3": ""}}, + clear=lambda: None, close=lambda: None, ) - with patch.object(self.service, "ModelCache", return_value=fake_models): - state = self.service.ServiceState("token", "instance") - state.register_client("client", "Blender") + with patch.object(self.service, "ModelManager", return_value=fake_models): + server = self.service.create_server(jobs_directory=self.jobs_directory) with tempfile.TemporaryDirectory() as directory: root = Path(directory) input_directory = root / "frames" input_directory.mkdir() (input_directory / "frame-0001.jpg").touch() (input_directory / "frame-0002.jpg").touch() - result_directory = root / "result" - status_path = root / "status.json" - state.reserve( - "job", - "client", + context = server.submit( "extract-frames", - str(status_path), - ) - - state.run_job( { - "job_id": "job", - "command": "extract-frames", - "parameters": { - "input": str(input_directory), - "output_dir": str(result_directory), - "status": str(status_path), - "video_frames": 16, - "start_frame": 1, - }, - } - ) - status = self.service.json.loads( - status_path.read_text(encoding="utf-8") + "input": str(input_directory), + "video_frames": 16, + "start_frame": 1, + }, + job_id="job", ) + server.run(context) + status = context.snapshot() self.assertTrue( - (result_directory / "frame.0001.jpg").is_file() + (context.directory / "frame.0001.jpg").is_file() ) self.assertEqual(status["state"], "succeeded") diff --git a/tests/test_worker.py b/tests/test_worker.py index e0ef592..5f94329 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -19,7 +19,7 @@ PROJECT_ROOT = Path(__file__).resolve().parents[1] def load_worker(): specification = importlib.util.spec_from_file_location( "snapform_worker", - PROJECT_ROOT / "src" / "snapform" / "worker.py", + PROJECT_ROOT / "src" / "snapform" / "ai_service" / "jobs.py", ) module = importlib.util.module_from_spec(specification) specification.loader.exec_module(module) @@ -500,18 +500,18 @@ class WorkerTest(unittest.TestCase): def test_models_from_r2_roundtrip_preserves_fields(self): root = PROJECT_ROOT / ".tmp-models-from-r2" model_dir = root / "models" / "DA3" - discovery = root / "service.json" + runtime_state = root / "runtime.json" try: - discovery.parent.mkdir(parents=True, exist_ok=True) - discovery.write_text( - json.dumps({"port": 1, "token": "t"}), + 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(discovery.read_text(encoding="utf-8")) - self.assertEqual(payload["port"], 1) + 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) diff --git a/tools/validate_depth_plane_asset.py b/tools/validate_depth_plane_asset.py index 967fb0e..669e3f6 100644 --- a/tools/validate_depth_plane_asset.py +++ b/tools/validate_depth_plane_asset.py @@ -6,7 +6,7 @@ import bpy sys.path.insert(0, str(Path(__file__).parents[1] / "src")) -from snapform.image_depth import ( +from snapform.blender.depth_plane import ( DEPTH_PLANE_MATERIAL_NODE_GROUP_NAME, SURFACE_MATERIAL_NODE_GROUP_NAME, add_depth_plane_processing, @@ -17,7 +17,7 @@ from snapform.image_depth import ( modifier_input_identifier, set_modifier_input, ) -from snapform.image_mesh import load_geometry_node_group +from snapform.blender.surface import load_geometry_node_group def main(): diff --git a/tools/validate_geometry_nodes_asset.py b/tools/validate_geometry_nodes_asset.py index 179457a..6ff650d 100644 --- a/tools/validate_geometry_nodes_asset.py +++ b/tools/validate_geometry_nodes_asset.py @@ -7,11 +7,11 @@ import numpy as np sys.path.insert(0, str(Path(__file__).parents[1] / "src")) -from snapform.image_depth import ( +from snapform.blender.depth_plane import ( get_modifier_input, set_modifier_input, ) -from snapform.image_mesh import create_matching_camera, load_geometry_node_group +from snapform.blender.surface import create_matching_camera, load_geometry_node_group def set_input(modifier, group, name, value): diff --git a/tools/validate_moge_depth_plane.py b/tools/validate_moge_depth_plane.py index 487f8db..7ab0427 100644 --- a/tools/validate_moge_depth_plane.py +++ b/tools/validate_moge_depth_plane.py @@ -9,7 +9,7 @@ import numpy as np PROJECT_ROOT = Path(__file__).parents[1] sys.path.insert(0, str(PROJECT_ROOT / "src")) -from snapform.image_depth import ( +from snapform.blender.depth_plane import ( MOGE_DEPTH_PLANE_NODE_GROUP_NAME, MOGE_MATERIAL_NODE_GROUP_NAME, build_moge2_depth_plane, diff --git a/tools/validate_moge_surface.py b/tools/validate_moge_surface.py index a0fce15..b25fba1 100644 --- a/tools/validate_moge_surface.py +++ b/tools/validate_moge_surface.py @@ -9,8 +9,8 @@ import numpy as np PROJECT_ROOT = Path(__file__).parents[1] sys.path.insert(0, str(PROJECT_ROOT / "src")) -from snapform.image_depth import MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME -from snapform.image_mesh import ( +from snapform.blender.depth_plane import MOGE_SURFACE_MATERIAL_NODE_GROUP_NAME +from snapform.blender.surface import ( MOGE_NODE_GROUP_NAME, build_image_mesh, prepare_moge_surface_archive,