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.
This commit is contained in:
@@ -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)。
|
||||
|
||||
179
docs/blendjob.md
Normal file
179
docs/blendjob.md
Normal file
@@ -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，并创建 `<jobs_directory>/<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 扩展增加安装依赖。
|
||||
@@ -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/
|
||||
```
|
||||
|
||||
|
||||
@@ -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`。
|
||||
|
||||
@@ -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)：开发验证、打包、上游版本与许可证
|
||||
|
||||
@@ -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/<repository-name>/<relative-path>`，并把 `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/<repository-name>/<relative-path>`。镜像选择写入持久的 `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/<opaque-job-id>/`，提交响应同时返回 `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`。
|
||||
|
||||
|
||||
@@ -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
|
||||
<opaque-job-id>/
|
||||
├── 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
|
||||
<opaque-job-id>/
|
||||
├── foreground.png 单图结果
|
||||
├── foreground.0001.png… 视频/序列帧结果
|
||||
├── frame.0001.jpg… 视频/序列帧的源帧副本
|
||||
|
||||
@@ -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,
|
||||
|
||||
1
src/snapform/ai_service/__init__.py
Normal file
1
src/snapform/ai_service/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""SnapForm AI Service implementation without Blender dependencies."""
|
||||
@@ -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/<name> 反推 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),
|
||||
)
|
||||
|
||||
|
||||
153
src/snapform/ai_service/model_manager.py
Normal file
153
src/snapform/ai_service/model_manager.py
Normal file
@@ -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()
|
||||
145
src/snapform/ai_service/server.py
Normal file
145
src/snapform/ai_service/server.py
Normal file
@@ -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()
|
||||
1
src/snapform/blender/__init__.py
Normal file
1
src/snapform/blender/__init__.py
Normal file
@@ -0,0 +1 @@
|
||||
"""Blender-process geometry construction and import logic."""
|
||||
@@ -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",
|
||||
@@ -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)
|
||||
37
src/snapform/blendjob/__init__.py
Normal file
37
src/snapform/blendjob/__init__.py
Normal file
@@ -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",
|
||||
)
|
||||
99
src/snapform/blendjob/client.py
Normal file
99
src/snapform/blendjob/client.py
Normal file
@@ -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,
|
||||
}
|
||||
267
src/snapform/blendjob/operator.py
Normal file
267
src/snapform/blendjob/operator.py
Normal file
@@ -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)
|
||||
350
src/snapform/blendjob/server.py
Normal file
350
src/snapform/blendjob/server.py
Normal file
@@ -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)
|
||||
208
src/snapform/blendjob/service.py
Normal file
208
src/snapform/blendjob/service.py
Normal file
@@ -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
|
||||
@@ -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):
|
||||
|
||||
@@ -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]
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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"
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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():
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user