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:
66
README.md
66
README.md
@@ -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
|
||||
|
||||
@@ -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
572
src/snapform/ai_host.py
Normal 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()
|
||||
@@ -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]
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,6 +1060,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
layout.prop(self, "snapform_depth_source")
|
||||
if self.snapform_depth_source == "ESTIMATE":
|
||||
estimate_box = layout.box()
|
||||
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(
|
||||
@@ -1049,10 +1076,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
||||
"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")
|
||||
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):
|
||||
|
||||
@@ -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,6 +80,7 @@ def close_active_job():
|
||||
|
||||
if _active_job is None:
|
||||
return
|
||||
if _active_job.log_file is not None:
|
||||
_active_job.log_file.close()
|
||||
if _active_job.temporary_input_path is not None:
|
||||
try:
|
||||
@@ -171,8 +178,18 @@ 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:
|
||||
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,
|
||||
@@ -180,7 +197,8 @@ class RunJob(bpy.types.Operator):
|
||||
cwd=runtime.storage_root(),
|
||||
creationflags=runtime.creation_flags(),
|
||||
)
|
||||
except OSError as error:
|
||||
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,6 +418,7 @@ class RunJob(bpy.types.Operator):
|
||||
def _finish(self, context, job, return_code):
|
||||
settings = context.scene.snapform_settings
|
||||
self._remove_timer(context)
|
||||
if job.log_file is not None:
|
||||
job.log_file.close()
|
||||
|
||||
if job.cancelled:
|
||||
@@ -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"}
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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"
|
||||
|
||||
@@ -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 HostJobRequest(
|
||||
{
|
||||
"job_id": Path(status_path).stem,
|
||||
"command": "reconstruct",
|
||||
"parameters": parameters,
|
||||
}
|
||||
)
|
||||
return command
|
||||
|
||||
|
||||
def creation_flags():
|
||||
|
||||
@@ -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")
|
||||
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
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
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}")
|
||||
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,
|
||||
|
||||
@@ -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"
|
||||
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
223
tests/test_ai_host.py
Normal 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()
|
||||
@@ -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__":
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user