Implement AI Host functionality with model caching and job management

- Added ai_host.py to handle AI tasks with FastAPI integration.
- Introduced HostState and ModelCache classes for managing client connections and model loading.
- Implemented job submission and cancellation endpoints in the AI host.
- Enhanced tests for AI host functionality, including client tracking and job cancellation.
- Updated existing tests to accommodate changes in the operator and command structure.
- Refactored runtime commands to use request-based architecture instead of command-line arguments.
- Improved model loading logic to handle CUDA memory management and reuse models efficiently.
This commit is contained in:
2026-08-04 07:35:40 +08:00
parent d694980da7
commit 838af0ee96
16 changed files with 2112 additions and 225 deletions

View File

@@ -9,15 +9,18 @@
## 当前功能
- 使用 Blender Extension manifest,可通过 `Install from Disk` 安装
- 在独立 Python Runtime 中运行 Torch 与 DA3,避免污染 Blender Python 进程
- 在独立 Python Runtime 的常驻 HTTP AI Host 中运行 Torch、DA3 与 BEN2,避免污染 Blender Python 进程
- Blender 使用内置 `urllib.request` 访问仅监听 `127.0.0.1` 的 Host;请求带随机 Token
- 常驻 Host 复用 Torch、CUDA、BEN2 与最近使用的 DA3 模型,连续执行时不再重复加载
- 多个 Blender 实例自动连接同一个共享 Host;`Host` 面板显示状态、端口、客户端与模型缓存,并提供安装、启动、清理、重启和停止控制
- Runtime、模型和任务结果默认保存在 `~/.snapform`,可由多个 Blender 版本共享
- 支持 DA3 Small、Base、Large 1.1、Giant 1.1 和 Nested Giant Large 1.1
- 支持 BEN2 Base 单图背景剔除测试,输出透明 PNG 与推理耗时
- 支持单张 JPG、JPEG、PNG、WEBP 图片或常见视频文件输入
- 视频均匀抽帧,可设置参与重建的最大帧数
- `Setup` 集中管理 Runtime,并逐项列出和下载所有可用模型
- `Host > Models` 默认折叠,展开后可逐项查看和下载所有可用模型
- Runtime 已就绪时不额外占用面板空间;模型列表每个模型只占一行
- `Debug` 是独立且默认折叠的 AI 功能测试面板
- `Debug` 默认折叠,内部的 DA3 与 BEN2 测试项也分别默认折叠
- 使用一个导出类型下拉菜单和一个 `Export` 按钮
- 只显示当前格式的专属参数;GLB 与 PLY 设置不会同时出现
- Debug 的模型选项按格式过滤;PLY 只提供两个带 Gaussian head 的 Giant 模型
@@ -31,7 +34,7 @@
- 转换对话框可选择推理模式、处理分辨率和面片网格细分
- Debug 的 DA3 导出或 BEN2 背景剔除完成后只打开对应 job 文件夹,不向当前 Blender 文件添加任何数据
- 完整保留每次任务生成的所有官方文件
- 安装、模型下载和推理均在后台子进程运行,状态栏实时显示阶段、百分比进度条与取消按钮
- 安装和模型下载使用后台子进程;推理提交给常驻 AI Host,状态栏实时显示阶段、百分比进度条与取消按钮
- 后台任务运行时,Install、Download 和 Export 按钮保持原位置并变灰,任务完成后恢复
## 使用
@@ -39,19 +42,20 @@
1. 使用 `uv run python pack.py` 构建 Extension ZIP。
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
4. 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3 和 BEN2 依赖。
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、DA3、BEN2、FastAPI、Uvicorn,
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
6. 在 `Setup` 中找到需要的模型并点击该行的 `Download`。
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
8. 选择一张图片或一个视频文件。视频可设置均匀抽取的帧数,然后点击 `Export`。
9. 测试背景剔除时,在 `Setup` 下载 `BEN2 Base`,然后在 `Debug` 的
9. 测试背景剔除时,确认 `Host > Models` 中的 `BEN2 Base` 已就绪,然后在 `Debug` 的
`BEN2 Background Removal` 区域选择单张图片并点击 `Remove Background`。
完成后会打开 job 文件夹,其中的 `foreground.png` 带透明 Alpha;状态栏显示
BEN2 推理、模型加载与任务总耗时。当前测试固定关闭较慢的 foreground refinement。
## Image Empty 去背景
1. 在 `Setup` 下载 `BEN2 Base`。
1. 确认 `Host > Models` 中的 `BEN2 Base` 已就绪;首次安装 Runtime 时会默认下载。
2. 在 3D View 选中 Image Empty,右键选择 `Remove Background`。
3. 完成后以 Fake User 保留原 Image datablock,新建并 Pack 一张透明 PNG 到当前 `.blend`,然后只把
当前 Empty 的 `data` 换成新图;Empty 的 Transform、显示尺寸、Offset 及其他设置保持不变。
@@ -66,16 +70,17 @@
1. 在 3D View 选中 `Object Type = Empty`、`Display As = Image` 的对象。
2. 右键选择 `Convert to Depth Plane`,通过 `Depth from` 选择深度来源:
- `AI Generated`:显示 Model、Use Ray Pose 和 Resolution Scale。实际 Process Res 为图片长边像素数
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `0.5`。计算后的 Process Res 会直接显示在
乘以 Resolution Scale;范围为 `0.1–1.0`,默认 `1.0`。计算后的 Process Res 会直接显示在
`Resolution Scale (… px)` 标签中。
转换导出 `mini_npz`,使用前需要在
`Setup` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
继续使用场景设置。
`Host > Models` 下载所选模型,并固定使用 `process_res_method="upper_bound_resize"`;设备和参考视图策略
继续使用场景设置。Model 下拉菜单只列出已经下载的 DA3 模型;若没有可用模型,对话框会引导到
`Host > Models` 下载。
- `Original`:直接把原图作为深度贴图,不显示估算参数。Remove Background 关闭时不启动
Runtime,也不要求源图片存在于磁盘。
`Original` 是默认深度来源。选择 `AI Generated` 时,其 Model、Use Ray Pose 与 Resolution Scale
参数会集中显示在独立边框中。两种模式始终显示 Remove Background、Displace via、Unshaded、
Subdivision 与 Delete Original。
参数会集中显示在独立边框中。两种模式都有默认折叠的 Options,内部依次为 Displace via、
Subdivision、Remove Background 和 Unshaded;Delete Original 保持在折叠项外。
Remove Background 默认关闭,与 Depth from 无关;开启后需要下载 `BEN2 Base`,并从原图生成、
Pack `.rb` 透明图片。它只替换生成面片材质的颜色图片,不修改源 Image Empty 的 Image datablock。
AI Generated 仍从原图生成 `.depth`;例如 `img.png` 会得到 `img.rb` 与 `img.depth`,不会生成
@@ -130,8 +135,9 @@ Principled BSDF 的 `Base Color`,并通过
`0.001`,会乘以面片 Transform 的统一缩放倍率。原图 Color 还会经过 `Float Curve` 连接到
Principled BSDF 的 `Roughness`,曲线两端输出分别为 `0.5` 和 `0.1`;材质 IOR 保持
Principled BSDF 默认值 `1.5`。
开启 `Unshaded` 后,Surface 改为原图 Color 直接连接 Emission,再连接 Material Output;不创建
Principled BSDF、Bump 或 Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
开启 `Unshaded` 后,原图 Color 连接 Emission,原图 Alpha 作为 Mix Shader 系数,在
Transparent BSDF 与 Emission 之间混合后连接 Material Output;不创建 Principled BSDF、Bump 或
Float Curve。若 Displace via 选择 `Material`,深度贴图与 Displacement
节点仍会保留并连接 Material Output 的 Displacement。
这里的“原始距离”是 DA3 模型输出的尺度。普通 Small、Base、Large 和 Giant 模型的深度不应
@@ -219,6 +225,7 @@ Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Pyt
- xFormers 0.0.33.post2(与 PyTorch 2.9.1 对应)
- 固定到已验证 commit 的 Depth Anything 3
- 固定到已验证 commit 的 BEN2
- FastAPI 0.139.2 与 Uvicorn 0.51.0
- 固定到 DA3 官方推荐 commit 的 `gsplat`;仅安装 PLY 导出所需的纯 Python 部分,
不编译本机 CUDA rasterizer
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
@@ -226,19 +233,43 @@ Extension 本身不包含 Torch。`Install Runtime` 使用 `uv` 创建独立 Pyt
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
通常很慢。
Extension 启用后只会查找已有的共享 AI Host,不会主动启动新进程。第一次执行 AI 操作或在侧边栏
手动点击 `Start` 时才会启动 Host。Host 只绑定回环地址,并让操作系统分配动态端口;第一个实例把
端口、进程和随机 Token 写入 `host.json`,其余 Blender 通过标准库
`urllib.request` 连接同一个 Host。`host.lock` 保证同一 Storage Root 同时只有一个 Host,Token
则阻止同机其他进程直接调用接口。Host 全局同时只执行一个 AI 任务。
每个 Blender 会注册独立客户端并每 5 秒发送心跳。正常关闭 Blender 时仅注销自身,不会直接终止
其他 Blender 正在使用的 Host;异常退出的客户端会在 20 秒后过期。最后一个客户端离开后,Host
保留模型缓存等待 60 秒,再在空闲时自动退出。取消任务不会重启 Host,并会保留已经加载的模型;
如果取消发生在一次不可中断的 Torch 调用中,任务会在该调用返回后的最近检查点结束。
侧边栏 `Host` 面板以 `Running (Clients: n)` 表示空闲服务,以 `Running (Busy)` 表示正在执行任务,
并显示缓存模型和活动任务。空闲且只有当前 Blender 连接时,可以清理模型缓存、重启或停止 Host;
多客户端或任务执行期间不会允许破坏共享服务。
BEN2 在首次使用后保持加载;DA3 保留最近使用的一个模型。切换 DA3 模型时会释放旧模型;加载时如果
显存不足,会先释放另一类缓存模型并重试。
Runtime Schema 已升级,旧 Runtime 需要重新执行 `Install Runtime` 才能安装 HTTP Host 依赖。
Host 的启动和崩溃日志写入 `~/.snapform/ai-host.log`。
Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`:
```text
~/.snapform/
├── runtime/
├── models/
└── jobs/
├── jobs/
├── ai-host.log
├── host.json # Host 运行时存在
└── host.lock # Host 运行时持有
```
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
`Install Runtime`;推理任务和模型读取可以共享。
`Install Runtime`;推理 Host、任务队列和模型读取可以共享。
## 开发
@@ -255,6 +286,7 @@ uv run python pack.py
├── src/
│ └── snapform/
│ ├── blender_manifest.toml
│ ├── ai_host.py
│ ├── bootstrap_runtime.py
│ ├── models.py
│ ├── image_depth.py

View File

@@ -1,5 +1,6 @@
import bpy
from . import runtime
from .image_depth import (
ApplyRemovedBackground,
BuildDepthPlane,
@@ -10,11 +11,18 @@ from .image_depth import (
)
from .operators import (
CancelJob,
ControlHost,
RunJob,
cancel_active_job,
close_active_job,
)
from .panel import DebugPanel, RuntimeModelsPanel
from .panel import (
Ben2DebugPanel,
Da3DebugPanel,
DebugPanel,
HostPanel,
ModelsPanel,
)
from .preferences import SnapFormPreferences
from .properties import SnapFormSettings
from .status_bar import draw_status_bar
@@ -25,13 +33,17 @@ CLASSES = (
SnapFormSettings,
RunJob,
CancelJob,
ControlHost,
ApplyRemovedBackground,
RemoveImageBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
ConvertImageToDepthPlane,
RuntimeModelsPanel,
HostPanel,
ModelsPanel,
DebugPanel,
Da3DebugPanel,
Ben2DebugPanel,
)
@@ -43,11 +55,14 @@ def register():
)
bpy.types.STATUSBAR_HT_header.append(draw_status_bar)
bpy.types.VIEW3D_MT_object_context_menu.append(draw_image_context_menu)
runtime.schedule_host_start()
def unregister():
cancel_active_job()
close_active_job()
runtime.unschedule_host()
runtime.detach_host_client()
bpy.types.VIEW3D_MT_object_context_menu.remove(draw_image_context_menu)
bpy.types.STATUSBAR_HT_header.remove(draw_status_bar)
if hasattr(bpy.types.Scene, "snapform_settings"):

572
src/snapform/ai_host.py Normal file
View File

@@ -0,0 +1,572 @@
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
HOST_SCHEMA = 2
CLIENT_TIMEOUT_SECONDS = 20.0
IDLE_SHUTDOWN_SECONDS = 60.0
class HostLock:
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):
import torch
self.torch = torch
self.ben2_model = None
self.ben2_key = None
self.da3_model = None
self.da3_key = None
def get_ben2(self, model_dir, device):
key = (str(Path(model_dir).resolve()), device.type)
if self.ben2_model is not None and self.ben2_key == key:
return self.ben2_model, 0.0
self.release_ben2()
try:
return self._load_ben2(model_dir, device, key)
except self.torch.cuda.OutOfMemoryError:
self.release_da3()
return self._load_ben2(model_dir, device, key)
def _load_ben2(self, model_dir, device, key):
from ben2 import BEN_Base
started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
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):
key = (str(Path(model_dir).resolve()), device.type)
if self.da3_model is not None and self.da3_key == key:
return self.da3_model
self.release_da3()
try:
return self._load_da3(model_dir, device, key)
except self.torch.cuda.OutOfMemoryError:
self.release_ben2()
return self._load_da3(model_dir, device, key)
def _load_da3(self, model_dir, device, key):
from depth_anything_3.api import DepthAnything3
model = DepthAnything3.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
self.da3_model = model
self.da3_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 close(self):
self.ben2_model = None
self.da3_model = None
self.ben2_key = None
self.da3_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": (
f"da3_{Path(self.da3_key[0]).name}"
if self.da3_key is not None
else ""
),
}
def _release_cuda_cache(self):
gc.collect()
if self.torch.cuda.is_available():
self.torch.cuda.empty_cache()
class HostState:
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,
"host_schema": HOST_SCHEMA,
"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,
"AI Host 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,
)
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="SnapFormHostMonitor",
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 = HostState(token, instance_id)
@asynccontextmanager
async def lifespan(_app):
state.start_monitor()
yield
state.close()
app = FastAPI(
title="SnapForm AI Host",
docs_url=None,
redoc_url=None,
openapi_url=None,
lifespan=lifespan,
)
app.state.snapform_host_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 host 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 {"remove-background", "reconstruct"}:
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="AI Host 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="AI Host is busy")
return {"state": "cleared"}
@app.post("/shutdown", status_code=202)
def shutdown_host(
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="AI Host 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(
{
"host_schema": HOST_SCHEMA,
"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()
host_lock = HostLock(args.lock)
if not host_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_host_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)
host_lock.release()
if __name__ == "__main__":
main()

View File

@@ -12,7 +12,7 @@ platforms = ["windows-x64"]
license = ["SPDX:GPL-3.0-or-later"]
[permissions]
network = "Download DA3 runtime dependencies and model weights"
network = "Access the local AI Host and download runtime dependencies and model weights"
files = "Read source images and store runtime data"
[build]

View File

@@ -6,16 +6,26 @@ import sys
from pathlib import Path
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
from models import BEN2_MODEL, MODELS
TORCH_VERSION = "2.9.1"
TORCHVISION_VERSION = "0.24.1"
XFORMERS_VERSION = "0.0.33.post2"
ADDICT_VERSION = "2.4.0"
FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0"
GSPLAT_COMMIT = "0b4dddf04cb687367602c01196913cde6a743d70"
RUNTIME_SCHEMA = 3
RUNTIME_SCHEMA = 4
TORCH_INDEXES = {
"cu128": "https://download.pytorch.org/whl/cu128",
"cpu": "https://download.pytorch.org/whl/cpu",
}
DEFAULT_MODEL_KEYS = ("SMALL", "BASE", BEN2_MODEL.key)
def write_status(path, progress, message, error=None):
@@ -46,6 +56,33 @@ def environment_python(runtime_dir):
return Path(runtime_dir) / "bin" / "python"
def download_default_models(args, python):
models = {**MODELS, BEN2_MODEL.key: BEN2_MODEL}
download_script = (
"import json, sys; "
"from huggingface_hub import snapshot_download; "
"patterns = json.loads(sys.argv[3]); "
"snapshot_download(repo_id=sys.argv[1], local_dir=sys.argv[2], "
"allow_patterns=patterns or None)"
)
models_dir = Path(args.models_dir)
models_dir.mkdir(parents=True, exist_ok=True)
for index, model_key in enumerate(DEFAULT_MODEL_KEYS):
model = models[model_key]
progress = 0.90 + index * 0.03
write_status(args.status, progress, f"Downloading {model.label}")
run(
[
python,
"-c",
download_script,
model.repository,
models_dir / model.key.lower(),
json.dumps(model.download_patterns),
]
)
def install(args):
runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
@@ -129,6 +166,19 @@ def install(args):
]
)
write_status(args.status, 0.75, "Installing AI Host")
run(
[
args.uv,
"pip",
"install",
"--python",
python,
f"fastapi=={FASTAPI_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
]
)
write_status(args.status, 0.78, "Installing gsplat build tools")
run(
[
@@ -162,7 +212,7 @@ def install(args):
gsplat_environment,
)
write_status(args.status, 0.92, "Checking Runtime")
write_status(args.status, 0.88, "Checking Runtime")
run(
[
python,
@@ -171,12 +221,15 @@ def install(args):
"import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"import fastapi, uvicorn; "
"import gsplat; "
"print(torch.__version__, torch.cuda.is_available())"
),
]
)
download_default_models(args, python)
ready_path = runtime_dir / "snapform_ready.json"
ready_path.write_text(
json.dumps(
@@ -185,6 +238,8 @@ def install(args):
"python": "3.12",
"torch": TORCH_VERSION,
"xformers": XFORMERS_VERSION,
"fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION,
"backend": args.backend,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
@@ -202,6 +257,7 @@ def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--uv", required=True)
parser.add_argument("--runtime-dir", required=True)
parser.add_argument("--models-dir", required=True)
parser.add_argument("--backend", choices=tuple(TORCH_INDEXES), required=True)
parser.add_argument("--da3-commit", required=True)
parser.add_argument("--ben2-commit", required=True)

View File

@@ -26,8 +26,13 @@ IMAGE_NAME_SUFFIXES = (
def image_depth_model_items(owner, context):
from .models import model_enum_items
from .runtime import model_ready
return model_enum_items(owner, context)
return [
item
for item in model_enum_items(owner, context)
if model_ready(item[0])
]
_last_depth_plane_model = "BASE"
@@ -921,8 +926,13 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy
return is_image_empty(context.object) and active_job() is None
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
)
def execute(self, context):
from .operators import IMAGE_BACKGROUND_JOB
@@ -1001,7 +1011,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
snapform_process_res_factor: bpy.props.FloatProperty(
name="Resolution Scale",
description="Scale applied to the source image long-edge resolution",
default=0.5,
default=1.0,
min=0.1,
max=1.0,
step=10,
@@ -1019,18 +1029,30 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
min=0,
max=10,
)
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy
return is_image_empty(context.object) and active_job() is None
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
)
def invoke(self, context, _event):
global _last_depth_plane_model
if not self.properties.is_property_set("snapform_model"):
self.snapform_model = _last_depth_plane_model
available_keys = [
item[0] for item in image_depth_model_items(self, context)
]
if available_keys:
self.snapform_model = (
_last_depth_plane_model
if _last_depth_plane_model in available_keys
else available_keys[0]
)
return context.window_manager.invoke_props_dialog(self, width=340)
def draw(self, context):
@@ -1038,21 +1060,32 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
layout.prop(self, "snapform_depth_source")
if self.snapform_depth_source == "ESTIMATE":
estimate_box = layout.box()
estimate_box.prop(self, "snapform_model")
estimate_box.prop(self, "snapform_use_ray_pose")
process_res = scaled_process_resolution(
context.object.data,
self.snapform_process_res_factor,
)
estimate_box.prop(
self,
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_displacement_via")
layout.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_subdivision_level")
available_models = image_depth_model_items(self, context)
if not available_models:
estimate_box.label(text="No DA3 models installed", icon="ERROR")
estimate_box.label(text="Download one in Host > Models")
else:
estimate_box.prop(self, "snapform_model")
estimate_box.prop(self, "snapform_use_ray_pose")
process_res = scaled_process_resolution(
context.object.data,
self.snapform_process_res_factor,
)
estimate_box.prop(
self,
"snapform_process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
settings_header, settings_body = layout.panel(
"snapform_depth_plane_options",
default_closed=True,
)
settings_header.label(text="Options")
if settings_body is not None:
settings_body.prop(self, "snapform_displacement_via")
settings_body.prop(self, "snapform_subdivision_level")
settings_body.prop(self, "snapform_remove_background")
settings_body.prop(self, "snapform_unshaded")
layout.prop(self, "snapform_delete_original")
def execute(self, context):

View File

@@ -34,7 +34,7 @@ def open_job_directory(result_directory):
@dataclass
class ActiveJob:
kind: str
process: subprocess.Popen
process: subprocess.Popen | None
log_file: object
log_path: Path
status_path: Path
@@ -50,6 +50,7 @@ class ActiveJob:
remove_background: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
_active_job = None
@@ -65,7 +66,12 @@ def cancel_active_job():
if _active_job is None:
return
_active_job.cancelled = True
if _active_job.process.poll() is None:
if _active_job.uses_host:
try:
runtime.cancel_host_job(_active_job.status_path.stem)
except RuntimeError:
pass
elif _active_job.process.poll() is None:
_active_job.process.terminate()
@@ -74,7 +80,8 @@ def close_active_job():
if _active_job is None:
return
_active_job.log_file.close()
if _active_job.log_file is not None:
_active_job.log_file.close()
if _active_job.temporary_input_path is not None:
try:
_active_job.temporary_input_path.unlink(missing_ok=True)
@@ -171,17 +178,28 @@ class RunJob(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
log_file = paths["log"].open("w", encoding="utf-8")
uses_host = isinstance(command, runtime.HostJobRequest)
log_file = None
log_path = paths["log"]
try:
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=runtime.creation_flags(),
)
except OSError as error:
log_file.close()
if uses_host:
runtime.submit_host_job(command)
process = None
log_path = runtime.host_log_path()
else:
if self.job_type == "INSTALL_RUNTIME":
runtime.stop_host()
log_file = paths["log"].open("w", encoding="utf-8")
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=runtime.creation_flags(),
)
except (OSError, RuntimeError) as error:
if log_file is not None:
log_file.close()
self.report({"ERROR"}, f"Unable to start task: {error}")
return {"CANCELLED"}
@@ -189,7 +207,7 @@ class RunJob(bpy.types.Operator):
kind=self.job_type,
process=process,
log_file=log_file,
log_path=paths["log"],
log_path=log_path,
status_path=paths["status"],
result_directory=paths["result"],
export_format=(
@@ -215,6 +233,7 @@ class RunJob(bpy.types.Operator):
and self.snapform_delete_input
else None
),
uses_host=uses_host,
)
settings.status = self._starting_message(settings)
settings.progress = 0.0
@@ -274,7 +293,7 @@ class RunJob(bpy.types.Operator):
if not runtime.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_command(
runtime.remove_background_request(
settings,
input_path,
paths["status"],
@@ -311,7 +330,7 @@ class RunJob(bpy.types.Operator):
"PLY export requires a DA3 Giant or Nested Giant model"
)
return (
runtime.reconstruct_command(
runtime.reconstruct_request(
settings,
export_format,
paths["status"],
@@ -372,6 +391,16 @@ class RunJob(bpy.types.Operator):
return {"CANCELLED"}
self._update_status(context, job)
if job.uses_host:
status = runtime.read_status(job.status_path) or {}
state = status.get("state")
if state == "succeeded":
return self._finish(context, job, 0)
if state in {"failed", "cancelled"}:
return self._finish(context, job, 1)
if runtime.host_status().get("state") in {"ERROR", "STOPPED"}:
return self._finish(context, job, 1)
return {"RUNNING_MODAL"}
return_code = job.process.poll()
if return_code is None:
return {"RUNNING_MODAL"}
@@ -389,7 +418,8 @@ class RunJob(bpy.types.Operator):
def _finish(self, context, job, return_code):
settings = context.scene.snapform_settings
self._remove_timer(context)
job.log_file.close()
if job.log_file is not None:
job.log_file.close()
if job.cancelled:
settings.status = "Task cancelled"
@@ -403,9 +433,14 @@ class RunJob(bpy.types.Operator):
if return_code != 0 or error:
settings.status = "Task failed"
settings.progress = 0.0
fallback_error = (
"AI Host task failed"
if job.uses_host
else "Task subprocess failed"
)
self.report(
{"ERROR"},
f"{error or 'Task subprocess failed'}; log: {job.log_path}",
f"{error or fallback_error}; log: {job.log_path}",
)
close_active_job()
self._redraw_ui(context, force=True)
@@ -559,3 +594,46 @@ class CancelJob(bpy.types.Operator):
def execute(self, _context):
cancel_active_job()
return {"FINISHED"}
class ControlHost(bpy.types.Operator):
bl_idname = "snapform.control_host"
bl_label = "Control SnapForm AI Host"
action: bpy.props.EnumProperty(
items=(
("START", "Start", "Start or connect to the shared AI Host"),
("CLEAR", "Clear Models", "Release cached AI models"),
("RESTART", "Restart", "Restart the shared AI Host"),
("STOP", "Stop", "Stop the shared AI Host"),
("LOG", "Open Log", "Open the AI Host log"),
),
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, _context):
try:
if self.action == "START":
runtime.enable_host()
message = "AI Host is ready"
elif self.action == "CLEAR":
runtime.clear_host_models()
message = "AI model cache is cleared"
elif self.action == "RESTART":
runtime.restart_host()
message = "AI Host is restarting"
elif self.action == "STOP":
runtime.stop_host()
message = "AI Host is stopping"
else:
log_path = runtime.host_log_path()
target = log_path if log_path.is_file() else log_path.parent
result = bpy.ops.wm.path_open(filepath=str(target))
if "FINISHED" not in result:
raise RuntimeError("Unable to open the AI Host log")
return {"FINISHED"}
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
self.report({"INFO"}, message)
return {"FINISHED"}

View File

@@ -1,21 +1,22 @@
import bpy
from .models import BEN2_MODEL, MODELS
from .models import BEN2_MODEL, MODELS, selected_model_key
from .operators import active_job
from . import runtime
from .runtime import model_ready, runtime_ready
class RuntimeModelsPanel(bpy.types.Panel):
bl_label = "Setup"
bl_idname = "SNAPFORM_PT_runtime_models"
class HostPanel(bpy.types.Panel):
bl_label = "Host"
bl_idname = "SNAPFORM_PT_host"
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
bl_order = 0
def draw(self, _context):
layout = self.layout
running = active_job() is not None
running = active_job() is not None or runtime.host_busy()
if not runtime_ready():
layout.label(text="Runtime is not installed", icon="ERROR")
install_row = layout.row()
@@ -26,6 +27,101 @@ class RuntimeModelsPanel(bpy.types.Panel):
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
return
status = runtime.host_status()
state = status.get("state", "STOPPED")
icons = {
"READY": "RADIOBUT_ON",
"BUSY": "TIME",
"STARTING": "TIME",
"ERROR": "ERROR",
"UNAVAILABLE": "ERROR",
"STOPPED": "PAUSE",
}
row = layout.row(align=True)
if state in {"READY", "BUSY"}:
row.label(
text=(
f"Running (Clients: {status.get('clients', 0)})"
if state == "READY"
else "Running (Busy)"
),
icon=icons[state],
)
if state == "READY":
controls = row.row(align=True)
controls.enabled = status.get("clients", 0) <= 1
operator = controls.operator(
"snapform.control_host",
text="",
icon="FILE_REFRESH",
)
operator.action = "RESTART"
operator = controls.operator(
"snapform.control_host",
text="",
icon="PAUSE",
)
operator.action = "STOP"
elif state == "STOPPED":
operator = row.operator(
"snapform.control_host",
text="Start",
icon="PLAY",
)
operator.action = "START"
else:
row.label(
text=status.get("message", state.title()),
icon=icons.get(state, "QUESTION"),
)
if state == "ERROR":
operator = row.operator(
"snapform.control_host",
text="Start",
icon="PLAY",
)
operator.action = "START"
if state in {"READY", "BUSY"}:
cached = status.get("cached_models", {})
cached_names = [name for name in cached.values() if name]
cached_row = layout.row(align=True)
cached_row.label(
text=(
f"Cached: {', '.join(cached_names)}"
if cached_names
else "Cached: None"
)
)
clear = cached_row.row(align=True)
clear.enabled = state == "READY"
operator = clear.operator(
"snapform.control_host",
text="",
icon="TRASH",
)
operator.action = "CLEAR"
class ModelsPanel(bpy.types.Panel):
bl_label = "Models"
bl_idname = "SNAPFORM_PT_models"
bl_parent_id = HostPanel.bl_idname
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 0
@classmethod
def poll(cls, _context):
return runtime_ready()
def draw(self, _context):
layout = self.layout
running = active_job() is not None or runtime.host_busy()
for model in MODELS.values():
self._draw_model(layout, model, running)
@@ -55,6 +151,24 @@ class DebugPanel(bpy.types.Panel):
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 1
def draw(self, context):
layout = self.layout
settings = context.scene.snapform_settings
layout.prop(settings, "device")
class Da3DebugPanel(bpy.types.Panel):
bl_label = "DA3 Reconstruction"
bl_idname = "SNAPFORM_PT_debug_da3"
bl_parent_id = DebugPanel.bl_idname
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 0
def draw(self, context):
layout = self.layout
@@ -82,7 +196,6 @@ class DebugPanel(bpy.types.Panel):
)
):
layout.prop(settings, "video_frames")
layout.prop(settings, "device")
layout.prop(settings, "process_res")
layout.prop(settings, "process_res_method")
layout.prop(settings, "use_ray_pose")
@@ -100,7 +213,11 @@ class DebugPanel(bpy.types.Panel):
ply_box.prop(settings, "gs_views_interval")
export_row = layout.row()
export_row.enabled = active_job() is None
export_row.enabled = (
active_job() is None
and not runtime.host_busy()
and model_ready(selected_model_key(settings))
)
operator = export_row.operator(
"snapform.run_job",
text="Export",
@@ -108,11 +225,28 @@ class DebugPanel(bpy.types.Panel):
)
operator.job_type = "RECONSTRUCT"
background_box = layout.box()
background_box.label(text="BEN2 Background Removal")
background_box.prop(settings, "background_input_path")
background_row = background_box.row()
background_row.enabled = active_job() is None
class Ben2DebugPanel(bpy.types.Panel):
bl_label = "BEN2 Background Removal"
bl_idname = "SNAPFORM_PT_debug_ben2"
bl_parent_id = DebugPanel.bl_idname
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 1
def draw(self, context):
layout = self.layout
settings = context.scene.snapform_settings
layout.prop(settings, "background_input_path")
background_row = layout.row()
background_row.enabled = (
active_job() is None
and not runtime.host_busy()
and model_ready(BEN2_MODEL.key)
)
operator = background_row.operator(
"snapform.run_job",
text="Remove Background",

View File

@@ -53,3 +53,10 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
layout.separator()
operator = layout.operator(
"snapform.control_host",
text="Open Host Log",
icon="TEXT",
)
operator.action = "LOG"

View File

@@ -1,10 +1,16 @@
import json
import os
import secrets
import shutil
import socket
import subprocess
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
@@ -13,10 +19,40 @@ from .models import BEN2_MODEL, get_downloadable_model, selected_model_key
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
RUNTIME_SCHEMA = 3
RUNTIME_SCHEMA = 4
HOST_SCHEMA = 2
HOST_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
@dataclass(frozen=True)
class HostJobRequest:
payload: dict
@dataclass
class HostConnection:
port: int
token: str
pid: int
instance_id: str
@dataclass
class LaunchedHost:
process: subprocess.Popen
log_file: object
_host_connection = None
_launched_host = None
_registered_instance_id = None
_client_id = uuid.uuid4().hex
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_local_opener = build_opener(ProxyHandler({}))
def preferences():
addon = bpy.context.preferences.addons.get(__package__)
return addon.preferences if addon is not None else None
@@ -124,6 +160,8 @@ def install_runtime_command(status_path):
executable,
"--runtime-dir",
str(runtime_directory()),
"--models-dir",
str(models_directory()),
"--backend",
backend.lower(),
"--da3-commit",
@@ -135,6 +173,419 @@ def install_runtime_command(status_path):
]
def host_log_path():
return storage_root() / "ai-host.log"
def host_discovery_path():
return storage_root() / "host.json"
def host_lock_path():
return storage_root() / "host.lock"
def _available_host_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind((HOST_ADDRESS, 0))
return listener.getsockname()[1]
def _host_url(connection, path):
return f"http://{HOST_ADDRESS}:{connection.port}{path}"
def _host_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(
_host_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"AI Host returned HTTP {error.code}")
except (OSError, URLError) as error:
raise RuntimeError(f"Unable to reach AI Host: {error}")
def _read_host_discovery():
try:
payload = json.loads(host_discovery_path().read_text(encoding="utf-8"))
return HostConnection(
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_host(connection, timeout=0.3):
health = _host_http_request(
connection,
"GET",
"/health",
timeout=timeout,
)
if health.get("host_schema") != HOST_SCHEMA:
raise RuntimeError("AI Host protocol is out of date")
if health.get("instance_id") != connection.instance_id:
raise RuntimeError("AI Host discovery is stale")
return health
def _discover_live_host(timeout=0.3):
connection = _read_host_discovery()
if connection is None:
return None, None
try:
return connection, _check_host(connection, timeout=timeout)
except RuntimeError:
return None, None
def start_host():
global _launched_host
if _launched_host is not None and _launched_host.process.poll() is None:
return
_close_launched_host()
port = _available_host_port()
token = secrets.token_urlsafe(32)
instance_id = uuid.uuid4().hex
log_file = host_log_path().open("a", encoding="utf-8")
command = [
str(runtime_python()),
"-u",
str(Path(__file__).with_name("ai_host.py")),
"--host",
HOST_ADDRESS,
"--port",
str(port),
"--token",
token,
"--instance-id",
instance_id,
"--discovery",
str(host_discovery_path()),
"--lock",
str(host_lock_path()),
]
try:
process = subprocess.Popen(
command,
stdout=log_file,
stderr=subprocess.STDOUT,
cwd=storage_root(),
creationflags=creation_flags(),
)
except OSError:
log_file.close()
raise
_launched_host = LaunchedHost(process, log_file)
def ensure_host(timeout=60.0):
global _host_auto_start, _host_connection, _host_snapshot
_host_auto_start = True
deadline = time.monotonic() + timeout
last_launch = 0.0
starting = False
while time.monotonic() < deadline:
connection, _health = _discover_live_host(timeout=0.5)
if connection is not None:
_host_connection = connection
_register_host_client(connection)
health = _check_host(connection, timeout=0.5)
_host_snapshot = _format_host_snapshot(connection, health)
_redraw_host_ui()
return connection
if not starting:
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
starting = True
_close_finished_launched_host()
if (
_launched_host is None
and time.monotonic() - last_launch >= 1.0
):
start_host()
last_launch = time.monotonic()
time.sleep(0.05)
raise RuntimeError(f"AI Host did not become ready; log: {host_log_path()}")
def submit_host_job(request):
connection = ensure_host()
payload = dict(request.payload)
payload["client_id"] = _client_id
_host_http_request(
connection,
"POST",
"/jobs",
payload=payload,
timeout=5.0,
)
return connection
def stop_host():
global _host_auto_start, _host_connection, _host_snapshot
global _registered_instance_id
connection = _host_connection
if connection is None:
connection, _health = _discover_live_host(timeout=0.5)
if connection is None:
_close_launched_host(terminate=True)
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
return
_register_host_client(connection)
_host_http_request(
connection,
"POST",
"/shutdown",
payload={"client_id": _client_id},
timeout=2.0,
)
_host_auto_start = False
deadline = time.monotonic() + 10.0
while time.monotonic() < deadline:
try:
_check_host(connection, timeout=0.2)
except RuntimeError:
break
time.sleep(0.05)
else:
raise RuntimeError("AI Host did not stop")
_host_connection = None
_registered_instance_id = None
_close_finished_launched_host()
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
def restart_host():
global _host_auto_start, _host_connection, _registered_instance_id
global _host_snapshot
connection = ensure_host()
_host_http_request(
connection,
"POST",
"/shutdown",
payload={"client_id": _client_id},
timeout=2.0,
)
_host_auto_start = True
_host_connection = None
_registered_instance_id = None
_host_snapshot = {"state": "STARTING", "message": "Restarting"}
def enable_host():
global _host_auto_start, _host_connection, _host_snapshot
_host_auto_start = True
connection, health = _discover_live_host(timeout=0.2)
if connection is not None:
_host_connection = connection
_register_host_client(connection)
_host_snapshot = _format_host_snapshot(connection, health)
_redraw_host_ui()
return
start_host()
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
def clear_host_models():
connection = ensure_host()
return _host_http_request(
connection,
"POST",
"/models/clear",
timeout=5.0,
)
def cancel_host_job(job_id):
connection = ensure_host()
return _host_http_request(
connection,
"DELETE",
f"/jobs/{job_id}",
payload={"client_id": _client_id},
timeout=2.0,
)
def detach_host_client():
global _host_connection, _registered_instance_id
connection = _host_connection
if connection is not None and _registered_instance_id is not None:
try:
_host_http_request(
connection,
"DELETE",
f"/clients/{_client_id}",
timeout=0.5,
)
except RuntimeError:
pass
_host_connection = None
_registered_instance_id = None
_close_launched_host(terminate=False)
def schedule_host_start():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_host_timer):
return
timers.register(_service_host_timer, first_interval=1.0, persistent=True)
def unschedule_host():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_host_timer):
timers.unregister(_service_host_timer)
def host_status():
return dict(_host_snapshot)
def host_busy():
return _host_snapshot.get("state") == "BUSY"
def _service_host_timer():
global _host_connection, _host_snapshot, _registered_instance_id
if not runtime_ready():
_host_snapshot = {
"state": "UNAVAILABLE",
"message": "Runtime is not installed",
}
return 2.0
connection, health = _discover_live_host(timeout=0.2)
if connection is None:
_host_connection = None
_registered_instance_id = None
_close_finished_launched_host()
if not _host_auto_start:
if _host_snapshot.get("state") != "STOPPED":
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_redraw_host_ui()
return 2.0
try:
start_host()
except OSError as error:
_host_snapshot = {"state": "ERROR", "message": str(error)}
return 2.0
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
return 1.0
try:
_host_connection = connection
if _registered_instance_id == connection.instance_id:
try:
_host_http_request(
connection,
"POST",
f"/clients/{_client_id}/heartbeat",
timeout=0.5,
)
except RuntimeError:
_registered_instance_id = None
_register_host_client(connection)
health = _check_host(connection, timeout=0.5)
_host_snapshot = _format_host_snapshot(connection, health)
except RuntimeError as error:
_host_connection = None
_registered_instance_id = None
_host_snapshot = {"state": "ERROR", "message": str(error)}
_redraw_host_ui()
return 5.0
def _register_host_client(connection):
global _registered_instance_id
if _registered_instance_id == connection.instance_id:
return
label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()
_host_http_request(
connection,
"POST",
"/clients",
payload={"client_id": _client_id, "label": label},
timeout=2.0,
)
_registered_instance_id = connection.instance_id
def _format_host_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,
}
def _redraw_host_ui():
try:
from .status_bar import redraw_ui
redraw_ui(bpy.context, force=True)
except (AttributeError, RuntimeError):
pass
def _close_finished_launched_host():
if _launched_host is None or _launched_host.process.poll() is None:
return
_close_launched_host()
def _close_launched_host(terminate=False):
global _launched_host
launched = _launched_host
_launched_host = 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 = [
@@ -153,31 +604,29 @@ def download_model_command(model_key, status_path):
return command
def remove_background_command(
def remove_background_request(
settings,
input_path,
status_path,
result_directory,
model_key,
):
return [
str(runtime_python()),
str(Path(__file__).with_name("worker.py")),
"remove-background",
"--input",
str(input_path),
"--model-dir",
str(model_directory(model_key)),
"--output-dir",
str(result_directory),
"--status",
str(status_path),
"--device",
settings.device.lower(),
]
return HostJobRequest(
{
"job_id": Path(status_path).stem,
"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": settings.device.lower(),
},
}
)
def reconstruct_command(
def reconstruct_request(
settings,
export_format,
status_path,
@@ -201,53 +650,39 @@ def reconstruct_command(
snapform_selected_ray_pose = (
settings.use_ray_pose if use_ray_pose is None else use_ray_pose
)
command = [
str(runtime_python()),
str(Path(__file__).with_name("worker.py")),
"reconstruct",
"--input",
str(
parameters = {
"input": str(
Path(input_path)
if input_path is not None
else Path(bpy.path.abspath(settings.input_path))
),
"--model-dir",
str(model_directory(snapform_selected_key)),
"--output-dir",
str(result_directory),
"--export-format",
export_format,
"--status",
str(status_path),
"--process-res",
str(snapform_selected_process_res),
"--process-res-method",
snapform_selected_process_res_method,
"--conf-thresh-percentile",
str(settings.conf_thresh_percentile),
"--num-max-points",
str(settings.num_max_points),
"--show-cameras",
"1" if settings.show_cameras else "0",
"--gs-views-interval",
str(settings.gs_views_interval),
"--video-frames",
str(settings.video_frames),
"--device",
settings.device.lower(),
"--use-ray-pose",
"1" if snapform_selected_ray_pose else "0",
"--ref-view-strategy",
settings.ref_view_strategy,
]
"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,
"conf_thresh_percentile": settings.conf_thresh_percentile,
"num_max_points": settings.num_max_points,
"show_cameras": 1 if settings.show_cameras else 0,
"gs_views_interval": settings.gs_views_interval,
"video_frames": settings.video_frames,
"device": settings.device.lower(),
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
"ref_view_strategy": settings.ref_view_strategy,
"remove_background_model_dir": "",
}
if remove_background:
command.extend(
[
"--remove-background-model-dir",
str(model_directory(BEN2_MODEL.key)),
]
parameters["remove_background_model_dir"] = str(
model_directory(BEN2_MODEL.key)
)
return command
return HostJobRequest(
{
"job_id": Path(status_path).stem,
"command": "reconstruct",
"parameters": parameters,
}
)
def creation_flags():

View File

@@ -23,10 +23,27 @@ SCENE_FILENAME = "scene.glb"
EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz")
def write_status(path, progress, message, error=None):
class TaskCancelled(RuntimeError):
pass
def check_cancelled(cancel_check):
if cancel_check is not None and cancel_check():
raise TaskCancelled("Task cancelled")
def write_status(path, progress, message, error=None, state=None):
payload = {"progress": progress, "message": message}
if error:
payload["error"] = error
if state is None:
if error:
state = "failed"
elif progress >= 1.0:
state = "succeeded"
else:
state = "running"
payload["state"] = state
Path(path).write_text(
json.dumps(payload, ensure_ascii=False),
encoding="utf-8",
@@ -42,7 +59,13 @@ def collect_images(directory):
)
def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
def extract_video_frames(
video_path,
output_directory,
maximum_frames,
cv2,
cancel_check=None,
):
capture = cv2.VideoCapture(str(video_path))
if not capture.isOpened():
raise RuntimeError(f"Unable to open video: {video_path}")
@@ -66,6 +89,7 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
image_paths = []
try:
for output_index, frame_index in enumerate(frame_indices):
check_cancelled(cancel_check)
capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index))
success, frame = capture.read()
if not success:
@@ -82,7 +106,8 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2):
return image_paths
def prepare_input(input_path, video_frames, temporary_root):
def prepare_input(input_path, video_frames, temporary_root, cancel_check=None):
check_cancelled(cancel_check)
input_path = Path(input_path)
suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES:
@@ -102,6 +127,7 @@ def prepare_input(input_path, video_frames, temporary_root):
temporary_directory,
video_frames,
cv2,
cancel_check=cancel_check,
)
except Exception:
shutil.rmtree(temporary_directory, ignore_errors=True)
@@ -203,19 +229,33 @@ def select_device(requested, torch):
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
def infer_foreground(input_path, model_dir, device, torch):
from ben2 import BEN_Base
def infer_foreground(
input_path,
model_dir,
device,
torch,
model_cache=None,
cancel_check=None,
):
from PIL import Image
input_path = Path(input_path)
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
load_started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
load_ms = (time.perf_counter() - load_started) * 1000.0
owns_model = model_cache is None
if owns_model:
from ben2 import BEN_Base
load_started = time.perf_counter()
model = BEN_Base.from_pretrained(model_dir)
model = model.to(device=device)
model.eval()
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
if device.type == "cuda":
@@ -229,17 +269,23 @@ def infer_foreground(input_path, model_dir, device, torch):
if device.type == "cuda":
torch.cuda.synchronize()
inference_ms = (time.perf_counter() - inference_started) * 1000.0
del model
if device.type == "cuda":
torch.cuda.empty_cache()
if owns_model:
del model
if device.type == "cuda":
torch.cuda.empty_cache()
check_cancelled(cancel_check)
return foreground, load_ms, inference_ms
def remove_background(args):
def remove_background(args, model_cache=None, cancel_check=None):
total_started = time.perf_counter()
write_status(args.status, 0.02, "Starting BEN2 Runtime")
check_cancelled(cancel_check)
import torch
if model_cache is None:
import torch
else:
torch = model_cache.torch
output_directory = Path(args.output_dir)
output_directory.mkdir(parents=True, exist_ok=False)
@@ -252,8 +298,11 @@ def remove_background(args):
args.model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
check_cancelled(cancel_check)
write_status(args.status, 0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
@@ -264,6 +313,7 @@ def remove_background(args):
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
check_cancelled(cancel_check)
total_ms = (time.perf_counter() - total_started) * 1000.0
message = (
f"Background removed in {inference_ms:.0f} ms "
@@ -288,10 +338,16 @@ def exported_paths(output_directory, export_format):
return []
def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
def wait_for_exported_paths(
output_directory,
export_format,
timeout=30.0,
cancel_check=None,
):
# 上游 NPZ exporter 使用后台线程,因此需要等待文件实际落盘。
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
check_cancelled(cancel_check)
result_paths = exported_paths(output_directory, export_format)
if result_paths:
return result_paths
@@ -299,11 +355,15 @@ def wait_for_exported_paths(output_directory, export_format, timeout=30.0):
return []
def reconstruct(args):
def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting AI runtime")
check_cancelled(cancel_check)
import torch
from depth_anything_3.api import DepthAnything3
if model_cache is None:
import torch
from depth_anything_3.api import DepthAnything3
else:
torch = model_cache.torch
from depth_anything_3.utils.export import export
output_directory = Path(args.output_dir)
@@ -313,9 +373,11 @@ def reconstruct(args):
args.input,
args.video_frames,
output_directory.parent,
cancel_check=cancel_check,
)
try:
check_cancelled(cancel_check)
device = select_device(args.device, torch)
remove_background_model_dir = getattr(
args,
@@ -329,7 +391,10 @@ def reconstruct(args):
remove_background_model_dir,
device,
torch,
model_cache=model_cache,
cancel_check=cancel_check,
)
check_cancelled(cancel_check)
foreground.convert("RGBA").save(
output_directory / "foreground.png",
format="PNG",
@@ -342,10 +407,14 @@ def reconstruct(args):
f"model load {load_ms:.0f} ms"
)
write_status(args.status, 0.15, f"Loading model on {device.type}")
model = DepthAnything3.from_pretrained(args.model_dir)
model = model.to(device=device)
model.eval()
if model_cache is None:
model = DepthAnything3.from_pretrained(args.model_dir)
model = model.to(device=device)
model.eval()
else:
model = model_cache.get_da3(args.model_dir, device)
check_cancelled(cancel_check)
write_status(
args.status,
0.25,
@@ -372,6 +441,7 @@ def reconstruct(args):
use_ray_pose=bool(args.use_ray_pose),
ref_view_strategy=args.ref_view_strategy,
)
check_cancelled(cancel_check)
write_status(args.status, 0.88, f"Exporting {args.export_format}")
export(
prediction,
@@ -379,15 +449,18 @@ def reconstruct(args):
str(output_directory),
**export_kwargs,
)
check_cancelled(cancel_check)
result_paths = wait_for_exported_paths(
output_directory,
args.export_format,
cancel_check=cancel_check,
)
if not result_paths:
raise RuntimeError(
f"No {args.export_format} files were exported"
)
check_cancelled(cancel_check)
write_status(
args.status,
1.0,

View File

@@ -343,7 +343,7 @@ class AddonTest(unittest.TestCase):
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual(
operator.__annotations__["snapform_process_res_factor"]["default"],
0.5,
1.0,
)
self.assertTrue(
operator.__annotations__["snapform_delete_original"]["default"]
@@ -403,7 +403,7 @@ class AddonTest(unittest.TestCase):
'"snapform_process_res_factor",'
),
draw_source.index(
'layout.prop(self, "snapform_subdivision_level")'
'settings_body.prop(self, "snapform_subdivision_level")'
),
)
self.assertIn(
@@ -420,17 +420,30 @@ class AddonTest(unittest.TestCase):
draw_source,
)
self.assertIn(
'layout.prop(self, "snapform_remove_background")',
'default_closed=True',
draw_source,
)
self.assertIn('settings_header.label(text="Options")', draw_source)
self.assertIn('settings_header.label(text="Options")', draw_source)
self.assertNotIn("layout.separator", draw_source)
self.assertNotIn("layout.label", draw_source)
self.assertIn(
'layout.prop(self, "snapform_unshaded")',
draw_source,
setting_fields = (
"snapform_displacement_via",
"snapform_subdivision_level",
"snapform_remove_background",
"snapform_unshaded",
)
model_items = operator.__annotations__["snapform_model"]["items"]
model_items = model_items(None, None)
field_positions = [
draw_source.index(f'settings_body.prop(self, "{field}")')
for field in setting_fields
]
self.assertEqual(field_positions, sorted(field_positions))
model_items_callback = operator.__annotations__["snapform_model"]["items"]
with patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
):
model_items = model_items_callback(None, None)
self.assertEqual(
[item[0] for item in model_items],
[
@@ -459,11 +472,33 @@ class AddonTest(unittest.TestCase):
)
self.addon.image_depth._last_depth_plane_model = "GIANT_1_1"
result = operator.invoke(context, None)
with patch.object(
self.addon.runtime,
"model_ready",
side_effect=lambda key: key == "GIANT_1_1",
):
result = operator.invoke(context, None)
self.assertIs(result, operator)
self.assertEqual(operator.snapform_model, "GIANT_1_1")
def test_conversion_dialog_only_lists_installed_models(self):
items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[
"snapform_model"
]["items"]
with patch.object(
self.addon.runtime,
"model_ready",
side_effect=lambda key: key in {"SMALL", "BASE"},
):
items = items_callback(None, None)
self.assertEqual([item[0] for item in items], ["SMALL", "BASE"])
draw_source = inspect.getsource(self.addon.ConvertImageToDepthPlane.draw)
self.assertIn('text="No DA3 models installed"', draw_source)
self.assertIn('text="Download one in Host > Models"', draw_source)
self.assertNotIn("_dialog_requirements_ready", inspect.getsource(self.addon.ConvertImageToDepthPlane))
def test_depth_conversion_progress_is_provider_neutral(self):
starting_source = inspect.getsource(
self.addon.RunJob._starting_message
@@ -532,7 +567,11 @@ class AddonTest(unittest.TestCase):
self.addon.image_depth,
"save_packed_image",
return_value=packed_path,
) as save_packed_image:
) as save_packed_image, patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
):
result = operator.execute(context)
self.assertEqual(result, {"FINISHED"})
@@ -640,8 +679,10 @@ class AddonTest(unittest.TestCase):
return_value=True,
) as model_ready, patch.object(
self.addon.runtime,
"reconstruct_command",
return_value=["worker"],
"reconstruct_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as reconstruct:
operator.snapform_model_key = snapform_model
operator.snapform_use_ray_pose = snapform_ray_pose
@@ -650,7 +691,10 @@ class AddonTest(unittest.TestCase):
settings,
)
self.assertEqual(command, ["worker"])
self.assertIsInstance(
command,
self.addon.runtime.HostJobRequest,
)
self.assertEqual(prepared_paths, paths)
model_ready.assert_called_once_with(snapform_model)
self.assertEqual(reconstruct.call_args.args[1], "mini_npz")
@@ -713,7 +757,7 @@ class AddonTest(unittest.TestCase):
with self.assertRaisesRegex(RuntimeError, "BEN2 Base"):
operator._prepare_command(None, settings)
def test_image_depth_background_option_reaches_worker_command(self):
def test_image_depth_background_option_reaches_host_request(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
@@ -739,7 +783,7 @@ class AddonTest(unittest.TestCase):
"model_directory",
side_effect=lambda key: Path(key.lower()),
):
command = runtime.reconstruct_command(
request = runtime.reconstruct_request(
settings,
"mini_npz",
Path("status.json"),
@@ -747,8 +791,10 @@ class AddonTest(unittest.TestCase):
remove_background=True,
)
option_index = command.index("--remove-background-model-dir")
self.assertEqual(command[option_index + 1], str(Path("ben2_base")))
self.assertEqual(
request.payload["parameters"]["remove_background_model_dir"],
str(Path("ben2_base")),
)
def test_preferences_use_package_identifier(self):
self.assertEqual(
@@ -767,6 +813,25 @@ class AddonTest(unittest.TestCase):
str(runtime.DEFAULT_STORAGE_ROOT),
)
def test_ai_host_uses_authenticated_loopback_http(self):
runtime = importlib.import_module("snapform.runtime")
request_source = inspect.getsource(runtime._host_http_request)
start_source = inspect.getsource(runtime.start_host)
self.assertEqual(runtime.HOST_ADDRESS, "127.0.0.1")
self.assertIn('"X-SnapForm-Token"', request_source)
self.assertIn('"/health"', inspect.getsource(runtime._check_host))
self.assertIn('with_name("ai_host.py")', start_source)
self.assertIn("ProxyHandler({})", inspect.getsource(runtime))
self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path))
def test_host_jobs_complete_from_status_state(self):
source = inspect.getsource(self.addon.RunJob.modal)
self.assertIn("if job.uses_host:", source)
self.assertIn('state == "succeeded"', source)
self.assertIn('{"failed", "cancelled"}', source)
def test_process_res_method_defaults_to_upper_bound_resize(self):
process_res = self.addon.SnapFormSettings.__annotations__[
"process_res"
@@ -780,7 +845,7 @@ class AddonTest(unittest.TestCase):
self.assertEqual(process_res["default"], 504)
self.assertEqual(process_res_method["default"], "upper_bound_resize")
def test_reconstruct_command_preserves_da3_parameter_names(self):
def test_reconstruct_request_preserves_da3_parameter_names(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
@@ -806,29 +871,30 @@ class AddonTest(unittest.TestCase):
"model_directory",
return_value=Path("model"),
):
command = runtime.reconstruct_command(
request = runtime.reconstruct_request(
settings,
"gs_ply",
Path("status.json"),
Path("result"),
)
self.assertIn("--export-format", command)
self.assertIn("gs_ply", command)
parameters = request.payload["parameters"]
self.assertEqual(request.payload["command"], "reconstruct")
self.assertEqual(parameters["export_format"], "gs_ply")
for parameter in (
"--process-res",
"--process-res-method",
"--conf-thresh-percentile",
"--num-max-points",
"--show-cameras",
"--gs-views-interval",
"--use-ray-pose",
"--ref-view-strategy",
"process_res",
"process_res_method",
"conf_thresh_percentile",
"num_max_points",
"show_cameras",
"gs_views_interval",
"use_ray_pose",
"ref_view_strategy",
):
with self.subTest(parameter=parameter):
self.assertIn(parameter, command)
self.assertIn(parameter, parameters)
def test_reconstruct_command_accepts_image_depth_overrides(self):
def test_reconstruct_request_accepts_image_depth_overrides(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
@@ -854,7 +920,7 @@ class AddonTest(unittest.TestCase):
"model_directory",
return_value=Path("large-model"),
) as snapform_model_directory:
snapform_command = runtime.reconstruct_command(
request = runtime.reconstruct_request(
settings,
"mini_npz",
Path("status.json"),
@@ -866,15 +932,13 @@ class AddonTest(unittest.TestCase):
)
snapform_model_directory.assert_called_once_with("BASE")
snapform_process_res_index = snapform_command.index("--process-res")
self.assertEqual(snapform_command[snapform_process_res_index + 1], "252")
snapform_process_res_method_index = snapform_command.index("--process-res-method")
parameters = request.payload["parameters"]
self.assertEqual(parameters["process_res"], 252)
self.assertEqual(
snapform_command[snapform_process_res_method_index + 1],
parameters["process_res_method"],
"lower_bound_resize",
)
snapform_ray_pose_index = snapform_command.index("--use-ray-pose")
self.assertEqual(snapform_command[snapform_ray_pose_index + 1], "0")
self.assertEqual(parameters["use_ray_pose"], 0)
def test_export_format_property_includes_all_da3_formats(self):
export_format = self.addon.SnapFormSettings.__annotations__[
@@ -904,37 +968,139 @@ class AddonTest(unittest.TestCase):
"DEFAULT_CLOSED",
self.addon.DebugPanel.bl_options,
)
for panel in (self.addon.Da3DebugPanel, self.addon.Ben2DebugPanel):
self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname)
self.assertIn("DEFAULT_CLOSED", panel.bl_options)
def test_debug_panel_uses_one_export_button_and_dynamic_settings(self):
draw_source = inspect.getsource(
self.addon.DebugPanel.draw
)
debug_source = inspect.getsource(self.addon.DebugPanel.draw)
da3_source = inspect.getsource(self.addon.Da3DebugPanel.draw)
ben2_source = inspect.getsource(self.addon.Ben2DebugPanel.draw)
self.assertEqual(draw_source.count("export_row.operator("), 1)
self.assertIn('settings.export_format == "glb"', draw_source)
self.assertIn('settings.export_format == "gs_ply"', draw_source)
self.assertIn('"gs_model" if settings.export_format', draw_source)
self.assertIn("layout.prop(settings, model_property)", draw_source)
self.assertNotIn("model_row", draw_source)
self.assertIn("export_row.enabled = active_job() is None", draw_source)
self.assertNotIn('text="Ready"', draw_source)
self.assertNotIn('text="Missing"', draw_source)
self.assertNotIn("License:", draw_source)
self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source)
self.assertIn('text="BEN2 Background Removal"', draw_source)
self.assertIn('layout.prop(settings, "device")', debug_source)
self.assertEqual(da3_source.count("export_row.operator("), 1)
self.assertIn('settings.export_format == "glb"', da3_source)
self.assertIn('settings.export_format == "gs_ply"', da3_source)
self.assertIn('"gs_model" if settings.export_format', da3_source)
self.assertIn("layout.prop(settings, model_property)", da3_source)
self.assertNotIn("model_row", da3_source)
self.assertIn("and model_ready(selected_model_key(settings))", da3_source)
self.assertNotIn('text="Ready"', da3_source)
self.assertNotIn('text="Missing"', da3_source)
self.assertNotIn("License:", da3_source)
self.assertIn('operator.job_type = "RECONSTRUCT"', da3_source)
self.assertIn(
'background_box.prop(settings, "background_input_path")',
draw_source,
'layout.prop(settings, "background_input_path")',
ben2_source,
)
self.assertIn('text="Remove Background"', draw_source)
self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', draw_source)
self.assertIn('text="Remove Background"', ben2_source)
self.assertIn("and model_ready(BEN2_MODEL.key)", ben2_source)
self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', ben2_source)
def test_runtime_models_panel_assigns_each_download_target(self):
def test_host_panel_exposes_shared_host_state_and_controls(self):
self.assertIn(self.addon.HostPanel, self.addon.CLASSES)
self.assertIn(self.addon.ControlHost, self.addon.CLASSES)
source = inspect.getsource(self.addon.HostPanel.draw)
self.assertEqual(self.addon.HostPanel.bl_label, "Host")
self.assertIn('text="Runtime is not installed"', source)
self.assertIn('text="Install Runtime"', source)
self.assertIn("runtime.host_status()", source)
self.assertIn('f"Running (Clients:', source)
self.assertIn('else "Running (Busy)"', source)
self.assertIn('"READY": "RADIOBUT_ON"', source)
self.assertIn('"BUSY": "TIME"', source)
self.assertNotIn('"READY": "PLAY"', source)
self.assertNotIn('text=f"Task:', source)
self.assertIn('elif state == "STOPPED":', source)
self.assertNotIn("Blender Clients:", source)
self.assertNotIn("status['port']", source)
self.assertIn('icon="TRASH"', source)
self.assertIn('icon="FILE_REFRESH"', source)
self.assertIn('icon="PAUSE"', source)
self.assertNotIn('text="Open Log"', source)
preferences_source = inspect.getsource(self.addon.SnapFormPreferences.draw)
self.assertIn('text="Open Host Log"', preferences_source)
def test_host_start_control_does_not_wait_for_model_runtime(self):
source = inspect.getsource(self.addon.runtime.enable_host)
self.assertIn("start_host()", source)
self.assertNotIn("ensure_host()", source)
def test_ai_operation_restarts_a_manually_stopped_host(self):
runtime = self.addon.runtime
connection = runtime.HostConnection(51234, "token", 42, "host")
health = {
"host_schema": runtime.HOST_SCHEMA,
"instance_id": "host",
"busy": False,
"clients": 1,
"cached_models": {},
}
previous_auto_start = runtime._host_auto_start
previous_connection = runtime._host_connection
previous_snapshot = runtime._host_snapshot
runtime._host_auto_start = False
runtime._host_connection = None
runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"}
try:
with patch.object(
runtime,
"_discover_live_host",
return_value=(connection, health),
), patch.object(
runtime,
"_register_host_client",
), patch.object(
runtime,
"_check_host",
return_value=health,
), patch.object(runtime, "_redraw_host_ui"):
result = runtime.ensure_host(timeout=0.1)
self.assertIs(result, connection)
self.assertTrue(runtime._host_auto_start)
self.assertEqual(runtime.host_status()["state"], "READY")
finally:
runtime._host_auto_start = previous_auto_start
runtime._host_connection = previous_connection
runtime._host_snapshot = previous_snapshot
def test_blender_startup_does_not_start_ai_host(self):
runtime = self.addon.runtime
previous_auto_start = runtime._host_auto_start
previous_snapshot = runtime._host_snapshot
runtime._host_auto_start = False
runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"}
try:
with patch.object(
runtime,
"runtime_ready",
return_value=True,
), patch.object(
runtime,
"_discover_live_host",
return_value=(None, None),
), patch.object(
runtime,
"_close_finished_launched_host",
), patch.object(runtime, "start_host") as start_host:
interval = runtime._service_host_timer()
start_host.assert_not_called()
self.assertEqual(interval, 2.0)
self.assertEqual(runtime.host_status()["state"], "STOPPED")
finally:
runtime._host_auto_start = previous_auto_start
runtime._host_snapshot = previous_snapshot
def test_models_panel_is_collapsed_and_assigns_each_download_target(self):
draw_source = inspect.getsource(
self.addon.RuntimeModelsPanel.draw
self.addon.ModelsPanel.draw
)
model_source = inspect.getsource(
self.addon.RuntimeModelsPanel._draw_model
self.addon.ModelsPanel._draw_model
)
self.assertIn("for model in MODELS.values()", draw_source)
@@ -951,8 +1117,16 @@ class AddonTest(unittest.TestCase):
self.assertNotIn("Isolated Runtime ready", draw_source)
self.assertNotIn("License:", model_source)
self.assertEqual(
self.addon.RuntimeModelsPanel.bl_label,
"Setup",
self.addon.ModelsPanel.bl_label,
"Models",
)
self.assertEqual(
self.addon.ModelsPanel.bl_parent_id,
self.addon.HostPanel.bl_idname,
)
self.assertEqual(
self.addon.ModelsPanel.bl_options,
{"DEFAULT_CLOSED"},
)
def test_ben2_debug_job_uses_its_image_and_model(self):
@@ -987,15 +1161,17 @@ class AddonTest(unittest.TestCase):
return_value=True,
) as model_ready, patch.object(
self.addon.runtime,
"remove_background_command",
return_value=["worker"],
"remove_background_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as remove_background:
command, prepared_paths = operator._prepare_command(
None,
settings,
)
self.assertEqual(command, ["worker"])
self.assertIsInstance(command, self.addon.runtime.HostJobRequest)
self.assertEqual(prepared_paths, paths)
model_ready.assert_called_once_with("BEN2_BASE")
self.assertEqual(remove_background.call_args.args[1], input_path)
@@ -1030,15 +1206,17 @@ class AddonTest(unittest.TestCase):
return_value=True,
), patch.object(
self.addon.runtime,
"remove_background_command",
return_value=["worker"],
"remove_background_request",
return_value=self.addon.runtime.HostJobRequest(
{"job_id": "test"}
),
) as remove_background:
command, prepared_paths = operator._prepare_command(
None,
settings,
)
self.assertEqual(command, ["worker"])
self.assertIsInstance(command, self.addon.runtime.HostJobRequest)
self.assertEqual(prepared_paths, paths)
self.assertEqual(remove_background.call_args.args[1], input_path)

223
tests/test_ai_host.py Normal file
View File

@@ -0,0 +1,223 @@
import importlib.util
import inspect
import sys
import tempfile
import unittest
from pathlib import Path
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
def load_ai_host():
specification = importlib.util.spec_from_file_location(
"snapform_ai_host",
PROJECT_ROOT / "src" / "snapform" / "ai_host.py",
)
module = importlib.util.module_from_spec(specification)
specification.loader.exec_module(module)
return module
class FakeCuda:
class OutOfMemoryError(RuntimeError):
pass
@staticmethod
def is_available():
return False
@staticmethod
def empty_cache():
return None
class HostTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.host = load_ai_host()
def create_cache(self):
cache = self.host.ModelCache.__new__(self.host.ModelCache)
cache.torch = SimpleNamespace(cuda=FakeCuda())
cache.ben2_model = None
cache.ben2_key = None
cache.da3_model = None
cache.da3_key = None
return cache
def test_ben2_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
ben2 = ModuleType("ben2")
ben2.BEN_Base = FakeModel
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(sys.modules, {"ben2": ben2}):
first, _load_ms = cache.get_ben2("model", device)
second, cached_load_ms = cache.get_ben2("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["ben2"], "model")
self.assertEqual(cached_load_ms, 0.0)
def test_da3_model_is_reused_for_matching_runtime_key(self):
loads = []
class FakeModel:
@classmethod
def from_pretrained(cls, model_dir):
loads.append(model_dir)
return cls()
def to(self, device):
self.device = device
return self
def eval(self):
return None
package = ModuleType("depth_anything_3")
api = ModuleType("depth_anything_3.api")
api.DepthAnything3 = FakeModel
cache = self.create_cache()
device = SimpleNamespace(type="cuda")
with patch.dict(
sys.modules,
{
"depth_anything_3": package,
"depth_anything_3.api": api,
},
):
first = cache.get_da3("model", device)
second = cache.get_da3("model", device)
self.assertIs(first, second)
self.assertEqual(loads, ["model"])
self.assertEqual(cache.status()["da3"], "da3_model")
def test_http_host_has_health_jobs_auth_and_single_task_guard(self):
source = inspect.getsource(self.host.create_app)
self.assertIn('@app.get("/health")', source)
self.assertIn('@app.post("/jobs", status_code=202)', source)
self.assertIn("x_snapform_token", source)
self.assertIn("hmac.compare_digest", source)
self.assertIn("parameters[\"status\"]", source)
self.assertIn('@app.post("/clients", status_code=201)', source)
self.assertIn('@app.post("/clients/{client_id}/heartbeat")', source)
self.assertIn('@app.delete("/jobs/{job_id}", status_code=202)', source)
self.assertIn('@app.post("/models/clear")', source)
self.assertIn('@app.post("/shutdown", status_code=202)', source)
self.assertNotIn("os._exit", source)
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("first", "Blender A")
state.register_client("second", "Blender B")
self.assertEqual(state.snapshot()["clients"], 2)
self.assertEqual(state.request_shutdown("first"), "shared")
self.assertTrue(
state.reserve("job", "first", "reconstruct", "status.json")
)
self.assertFalse(
state.reserve("other", "second", "reconstruct", "other.json")
)
def test_cancelling_job_keeps_host_and_model_cache_alive(self):
fake_models = SimpleNamespace(
status=lambda: {"ben2": "BEN2_BASE", "da3": ""},
close=lambda: None,
)
with patch.object(self.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("client", "Blender")
with tempfile.TemporaryDirectory() as directory:
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.host.json.loads(
status_path.read_text(encoding="utf-8")
)
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.host, "ModelCache", return_value=fake_models):
state = self.host.HostState("token", "instance")
state.register_client("client", "Blender")
state.clients["client"]["last_seen"] -= (
self.host.CLIENT_TIMEOUT_SECONDS + 1.0
)
state.expire_clients()
self.assertEqual(state.snapshot()["clients"], 0)
def test_discovery_is_only_removed_by_matching_host(self):
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "host.json"
args = SimpleNamespace(
discovery=str(path),
port=51234,
token="token",
instance_id="owner",
)
with patch.object(self.host.os, "getpid", return_value=42):
self.host.write_discovery(args)
payload = self.host.json.loads(path.read_text(encoding="utf-8"))
self.assertEqual(payload["port"], 51234)
self.assertEqual(payload["pid"], 42)
self.host.remove_discovery(path, "other")
self.assertTrue(path.is_file())
self.host.remove_discovery(path, "owner")
self.assertFalse(path.exists())
if __name__ == "__main__":
unittest.main()

View File

@@ -29,6 +29,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
args = SimpleNamespace(
uv="uv",
runtime_dir=str(runtime_dir),
models_dir=str(root / "models"),
backend="cu128",
da3_commit="test-commit",
ben2_commit="test-ben2-commit",
@@ -67,13 +68,48 @@ class BootstrapRuntimeTest(unittest.TestCase):
self.assertTrue(
any("test-ben2-commit" in str(value) for value in ben2_command)
)
check_command = commands[-1]
host_command = next(
command
for command in commands
if f"fastapi=={bootstrap.FASTAPI_VERSION}" in command
)
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
check_command = next(
command
for command in commands
if any("import torch;" in str(value) for value in command)
)
self.assertTrue(
any(
"from ben2 import BEN_Base" in str(value)
for value in check_command
)
)
self.assertTrue(
any("import fastapi, uvicorn" in str(value) for value in check_command)
)
default_downloads = [
command
for command in commands
if any(
repository in command
for repository in (
"depth-anything/DA3-SMALL",
"depth-anything/DA3-BASE",
"PramaLLC/BEN2",
)
)
]
self.assertEqual(len(default_downloads), 3)
ben2_download = next(
command
for command in default_downloads
if "PramaLLC/BEN2" in command
)
self.assertEqual(
ben2_download[-1],
'["config.json", "model.safetensors"]',
)
if __name__ == "__main__":

View File

@@ -38,5 +38,6 @@ class PackTest(unittest.TestCase):
self.assertIn("image_depth.py", names)
self.assertNotIn("snapform_depth_plane.py", names)
self.assertIn("worker.py", names)
self.assertIn("ai_host.py", names)
self.assertNotIn("importer.py", names)
self.assertNotIn("snapform/__init__.py", names)

View File

@@ -47,6 +47,20 @@ class WorkerTest(unittest.TestCase):
source.index("import torch"),
)
def test_ai_tasks_support_cooperative_cancellation(self):
with self.assertRaises(self.worker.TaskCancelled):
self.worker.check_cancelled(lambda: True)
self.worker.check_cancelled(lambda: False)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.reconstruct),
)
self.assertIn(
"cancel_check=cancel_check",
inspect.getsource(self.worker.remove_background),
)
def test_optional_background_removal_runs_before_da3_model_load(self):
source = inspect.getsource(self.worker.reconstruct)