diff --git a/README.md b/README.md index 007867b..d05c296 100644 --- a/README.md +++ b/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 diff --git a/src/snapform/__init__.py b/src/snapform/__init__.py index d80124f..c77d301 100644 --- a/src/snapform/__init__.py +++ b/src/snapform/__init__.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"): diff --git a/src/snapform/ai_host.py b/src/snapform/ai_host.py new file mode 100644 index 0000000..bcd389c --- /dev/null +++ b/src/snapform/ai_host.py @@ -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() diff --git a/src/snapform/blender_manifest.toml b/src/snapform/blender_manifest.toml index 633be26..52b6448 100644 --- a/src/snapform/blender_manifest.toml +++ b/src/snapform/blender_manifest.toml @@ -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] diff --git a/src/snapform/bootstrap_runtime.py b/src/snapform/bootstrap_runtime.py index 1e91425..cd3e6fe 100644 --- a/src/snapform/bootstrap_runtime.py +++ b/src/snapform/bootstrap_runtime.py @@ -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) diff --git a/src/snapform/image_depth.py b/src/snapform/image_depth.py index 97494d1..d71ec2e 100644 --- a/src/snapform/image_depth.py +++ b/src/snapform/image_depth.py @@ -26,8 +26,13 @@ IMAGE_NAME_SUFFIXES = ( def image_depth_model_items(owner, context): from .models import model_enum_items + from .runtime import model_ready - return model_enum_items(owner, context) + return [ + item + for item in model_enum_items(owner, context) + if model_ready(item[0]) + ] _last_depth_plane_model = "BASE" @@ -921,8 +926,13 @@ class RemoveImageBackground(bpy.types.Operator): @classmethod def poll(cls, context): from .operators import active_job + from .runtime import host_busy - return is_image_empty(context.object) and active_job() is None + return ( + is_image_empty(context.object) + and active_job() is None + and not host_busy() + ) def execute(self, context): from .operators import IMAGE_BACKGROUND_JOB @@ -1001,7 +1011,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator): snapform_process_res_factor: bpy.props.FloatProperty( name="Resolution Scale", description="Scale applied to the source image long-edge resolution", - default=0.5, + default=1.0, min=0.1, max=1.0, step=10, @@ -1019,18 +1029,30 @@ class ConvertImageToDepthPlane(bpy.types.Operator): min=0, max=10, ) - @classmethod def poll(cls, context): from .operators import active_job + from .runtime import host_busy - return is_image_empty(context.object) and active_job() is None + return ( + is_image_empty(context.object) + and active_job() is None + and not host_busy() + ) def invoke(self, context, _event): global _last_depth_plane_model if not self.properties.is_property_set("snapform_model"): - self.snapform_model = _last_depth_plane_model + available_keys = [ + item[0] for item in image_depth_model_items(self, context) + ] + if available_keys: + self.snapform_model = ( + _last_depth_plane_model + if _last_depth_plane_model in available_keys + else available_keys[0] + ) return context.window_manager.invoke_props_dialog(self, width=340) def draw(self, context): @@ -1038,21 +1060,32 @@ class ConvertImageToDepthPlane(bpy.types.Operator): layout.prop(self, "snapform_depth_source") if self.snapform_depth_source == "ESTIMATE": estimate_box = layout.box() - estimate_box.prop(self, "snapform_model") - estimate_box.prop(self, "snapform_use_ray_pose") - process_res = scaled_process_resolution( - context.object.data, - self.snapform_process_res_factor, - ) - estimate_box.prop( - self, - "snapform_process_res_factor", - text=f"Resolution Scale ({process_res} px)", - ) - layout.prop(self, "snapform_remove_background") - layout.prop(self, "snapform_displacement_via") - layout.prop(self, "snapform_unshaded") - layout.prop(self, "snapform_subdivision_level") + available_models = image_depth_model_items(self, context) + if not available_models: + estimate_box.label(text="No DA3 models installed", icon="ERROR") + estimate_box.label(text="Download one in Host > Models") + else: + estimate_box.prop(self, "snapform_model") + estimate_box.prop(self, "snapform_use_ray_pose") + process_res = scaled_process_resolution( + context.object.data, + self.snapform_process_res_factor, + ) + estimate_box.prop( + self, + "snapform_process_res_factor", + text=f"Resolution Scale ({process_res} px)", + ) + settings_header, settings_body = layout.panel( + "snapform_depth_plane_options", + default_closed=True, + ) + settings_header.label(text="Options") + if settings_body is not None: + settings_body.prop(self, "snapform_displacement_via") + settings_body.prop(self, "snapform_subdivision_level") + settings_body.prop(self, "snapform_remove_background") + settings_body.prop(self, "snapform_unshaded") layout.prop(self, "snapform_delete_original") def execute(self, context): diff --git a/src/snapform/operators.py b/src/snapform/operators.py index 3c62c38..91ae12c 100644 --- a/src/snapform/operators.py +++ b/src/snapform/operators.py @@ -34,7 +34,7 @@ def open_job_directory(result_directory): @dataclass class ActiveJob: kind: str - process: subprocess.Popen + process: subprocess.Popen | None log_file: object log_path: Path status_path: Path @@ -50,6 +50,7 @@ class ActiveJob: remove_background: bool = False temporary_input_path: Path | None = None cancelled: bool = False + uses_host: bool = False _active_job = None @@ -65,7 +66,12 @@ def cancel_active_job(): if _active_job is None: return _active_job.cancelled = True - if _active_job.process.poll() is None: + if _active_job.uses_host: + try: + runtime.cancel_host_job(_active_job.status_path.stem) + except RuntimeError: + pass + elif _active_job.process.poll() is None: _active_job.process.terminate() @@ -74,7 +80,8 @@ def close_active_job(): if _active_job is None: return - _active_job.log_file.close() + if _active_job.log_file is not None: + _active_job.log_file.close() if _active_job.temporary_input_path is not None: try: _active_job.temporary_input_path.unlink(missing_ok=True) @@ -171,17 +178,28 @@ class RunJob(bpy.types.Operator): self.report({"ERROR"}, str(error)) return {"CANCELLED"} - log_file = paths["log"].open("w", encoding="utf-8") + uses_host = isinstance(command, runtime.HostJobRequest) + log_file = None + log_path = paths["log"] try: - process = subprocess.Popen( - command, - stdout=log_file, - stderr=subprocess.STDOUT, - cwd=runtime.storage_root(), - creationflags=runtime.creation_flags(), - ) - except OSError as error: - log_file.close() + if uses_host: + runtime.submit_host_job(command) + process = None + log_path = runtime.host_log_path() + else: + if self.job_type == "INSTALL_RUNTIME": + runtime.stop_host() + log_file = paths["log"].open("w", encoding="utf-8") + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + cwd=runtime.storage_root(), + creationflags=runtime.creation_flags(), + ) + except (OSError, RuntimeError) as error: + if log_file is not None: + log_file.close() self.report({"ERROR"}, f"Unable to start task: {error}") return {"CANCELLED"} @@ -189,7 +207,7 @@ class RunJob(bpy.types.Operator): kind=self.job_type, process=process, log_file=log_file, - log_path=paths["log"], + log_path=log_path, status_path=paths["status"], result_directory=paths["result"], export_format=( @@ -215,6 +233,7 @@ class RunJob(bpy.types.Operator): and self.snapform_delete_input else None ), + uses_host=uses_host, ) settings.status = self._starting_message(settings) settings.progress = 0.0 @@ -274,7 +293,7 @@ class RunJob(bpy.types.Operator): if not runtime.model_ready(BEN2_MODEL.key): raise RuntimeError(f"Download {BEN2_MODEL.label} first") return ( - runtime.remove_background_command( + runtime.remove_background_request( settings, input_path, paths["status"], @@ -311,7 +330,7 @@ class RunJob(bpy.types.Operator): "PLY export requires a DA3 Giant or Nested Giant model" ) return ( - runtime.reconstruct_command( + runtime.reconstruct_request( settings, export_format, paths["status"], @@ -372,6 +391,16 @@ class RunJob(bpy.types.Operator): return {"CANCELLED"} self._update_status(context, job) + if job.uses_host: + status = runtime.read_status(job.status_path) or {} + state = status.get("state") + if state == "succeeded": + return self._finish(context, job, 0) + if state in {"failed", "cancelled"}: + return self._finish(context, job, 1) + if runtime.host_status().get("state") in {"ERROR", "STOPPED"}: + return self._finish(context, job, 1) + return {"RUNNING_MODAL"} return_code = job.process.poll() if return_code is None: return {"RUNNING_MODAL"} @@ -389,7 +418,8 @@ class RunJob(bpy.types.Operator): def _finish(self, context, job, return_code): settings = context.scene.snapform_settings self._remove_timer(context) - job.log_file.close() + if job.log_file is not None: + job.log_file.close() if job.cancelled: settings.status = "Task cancelled" @@ -403,9 +433,14 @@ class RunJob(bpy.types.Operator): if return_code != 0 or error: settings.status = "Task failed" settings.progress = 0.0 + fallback_error = ( + "AI Host task failed" + if job.uses_host + else "Task subprocess failed" + ) self.report( {"ERROR"}, - f"{error or 'Task subprocess failed'}; log: {job.log_path}", + f"{error or fallback_error}; log: {job.log_path}", ) close_active_job() self._redraw_ui(context, force=True) @@ -559,3 +594,46 @@ class CancelJob(bpy.types.Operator): def execute(self, _context): cancel_active_job() return {"FINISHED"} + + +class ControlHost(bpy.types.Operator): + bl_idname = "snapform.control_host" + bl_label = "Control SnapForm AI Host" + + action: bpy.props.EnumProperty( + items=( + ("START", "Start", "Start or connect to the shared AI Host"), + ("CLEAR", "Clear Models", "Release cached AI models"), + ("RESTART", "Restart", "Restart the shared AI Host"), + ("STOP", "Stop", "Stop the shared AI Host"), + ("LOG", "Open Log", "Open the AI Host log"), + ), + options={"HIDDEN", "SKIP_SAVE"}, + ) + + def execute(self, _context): + try: + if self.action == "START": + runtime.enable_host() + message = "AI Host is ready" + elif self.action == "CLEAR": + runtime.clear_host_models() + message = "AI model cache is cleared" + elif self.action == "RESTART": + runtime.restart_host() + message = "AI Host is restarting" + elif self.action == "STOP": + runtime.stop_host() + message = "AI Host is stopping" + else: + log_path = runtime.host_log_path() + target = log_path if log_path.is_file() else log_path.parent + result = bpy.ops.wm.path_open(filepath=str(target)) + if "FINISHED" not in result: + raise RuntimeError("Unable to open the AI Host log") + return {"FINISHED"} + except (OSError, RuntimeError) as error: + self.report({"ERROR"}, str(error)) + return {"CANCELLED"} + self.report({"INFO"}, message) + return {"FINISHED"} diff --git a/src/snapform/panel.py b/src/snapform/panel.py index 65c61ce..bab7416 100644 --- a/src/snapform/panel.py +++ b/src/snapform/panel.py @@ -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", diff --git a/src/snapform/preferences.py b/src/snapform/preferences.py index e3e573d..346ce1c 100644 --- a/src/snapform/preferences.py +++ b/src/snapform/preferences.py @@ -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" diff --git a/src/snapform/runtime.py b/src/snapform/runtime.py index e608ff9..0a75b82 100644 --- a/src/snapform/runtime.py +++ b/src/snapform/runtime.py @@ -1,10 +1,16 @@ import json import os +import secrets import shutil +import socket import subprocess +import time import uuid +from dataclasses import dataclass from datetime import datetime from pathlib import Path +from urllib.error import HTTPError, URLError +from urllib.request import ProxyHandler, Request, build_opener import bpy @@ -13,10 +19,40 @@ from .models import BEN2_MODEL, get_downloadable_model, selected_model_key DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4" BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6" -RUNTIME_SCHEMA = 3 +RUNTIME_SCHEMA = 4 +HOST_SCHEMA = 2 +HOST_ADDRESS = "127.0.0.1" DEFAULT_STORAGE_ROOT = Path.home() / ".snapform" +@dataclass(frozen=True) +class HostJobRequest: + payload: dict + + +@dataclass +class HostConnection: + port: int + token: str + pid: int + instance_id: str + + +@dataclass +class LaunchedHost: + process: subprocess.Popen + log_file: object + + +_host_connection = None +_launched_host = None +_registered_instance_id = None +_client_id = uuid.uuid4().hex +_host_auto_start = False +_host_snapshot = {"state": "STOPPED", "message": "Stopped"} +_local_opener = build_opener(ProxyHandler({})) + + def preferences(): addon = bpy.context.preferences.addons.get(__package__) return addon.preferences if addon is not None else None @@ -124,6 +160,8 @@ def install_runtime_command(status_path): executable, "--runtime-dir", str(runtime_directory()), + "--models-dir", + str(models_directory()), "--backend", backend.lower(), "--da3-commit", @@ -135,6 +173,419 @@ def install_runtime_command(status_path): ] +def host_log_path(): + return storage_root() / "ai-host.log" + + +def host_discovery_path(): + return storage_root() / "host.json" + + +def host_lock_path(): + return storage_root() / "host.lock" + + +def _available_host_port(): + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener: + listener.bind((HOST_ADDRESS, 0)) + return listener.getsockname()[1] + + +def _host_url(connection, path): + return f"http://{HOST_ADDRESS}:{connection.port}{path}" + + +def _host_http_request(connection, method, path, payload=None, timeout=2.0): + data = None + headers = {"X-SnapForm-Token": connection.token} + if payload is not None: + data = json.dumps(payload).encode("utf-8") + headers["Content-Type"] = "application/json" + request = Request( + _host_url(connection, path), + data=data, + headers=headers, + method=method, + ) + try: + with _local_opener.open(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + except HTTPError as error: + try: + detail = json.loads(error.read().decode("utf-8")).get("detail") + except (json.JSONDecodeError, UnicodeDecodeError): + detail = None + raise RuntimeError(detail or f"AI Host returned HTTP {error.code}") + except (OSError, URLError) as error: + raise RuntimeError(f"Unable to reach AI Host: {error}") + + +def _read_host_discovery(): + try: + payload = json.loads(host_discovery_path().read_text(encoding="utf-8")) + return HostConnection( + port=int(payload["port"]), + token=str(payload["token"]), + pid=int(payload["pid"]), + instance_id=str(payload["instance_id"]), + ) + except (KeyError, TypeError, ValueError, OSError, json.JSONDecodeError): + return None + + +def _check_host(connection, timeout=0.3): + health = _host_http_request( + connection, + "GET", + "/health", + timeout=timeout, + ) + if health.get("host_schema") != HOST_SCHEMA: + raise RuntimeError("AI Host protocol is out of date") + if health.get("instance_id") != connection.instance_id: + raise RuntimeError("AI Host discovery is stale") + return health + + +def _discover_live_host(timeout=0.3): + connection = _read_host_discovery() + if connection is None: + return None, None + try: + return connection, _check_host(connection, timeout=timeout) + except RuntimeError: + return None, None + + +def start_host(): + global _launched_host + + if _launched_host is not None and _launched_host.process.poll() is None: + return + _close_launched_host() + port = _available_host_port() + token = secrets.token_urlsafe(32) + instance_id = uuid.uuid4().hex + log_file = host_log_path().open("a", encoding="utf-8") + command = [ + str(runtime_python()), + "-u", + str(Path(__file__).with_name("ai_host.py")), + "--host", + HOST_ADDRESS, + "--port", + str(port), + "--token", + token, + "--instance-id", + instance_id, + "--discovery", + str(host_discovery_path()), + "--lock", + str(host_lock_path()), + ] + try: + process = subprocess.Popen( + command, + stdout=log_file, + stderr=subprocess.STDOUT, + cwd=storage_root(), + creationflags=creation_flags(), + ) + except OSError: + log_file.close() + raise + _launched_host = LaunchedHost(process, log_file) + + +def ensure_host(timeout=60.0): + global _host_auto_start, _host_connection, _host_snapshot + + _host_auto_start = True + deadline = time.monotonic() + timeout + last_launch = 0.0 + starting = False + while time.monotonic() < deadline: + connection, _health = _discover_live_host(timeout=0.5) + if connection is not None: + _host_connection = connection + _register_host_client(connection) + health = _check_host(connection, timeout=0.5) + _host_snapshot = _format_host_snapshot(connection, health) + _redraw_host_ui() + return connection + if not starting: + _host_snapshot = {"state": "STARTING", "message": "Starting"} + _redraw_host_ui() + starting = True + _close_finished_launched_host() + if ( + _launched_host is None + and time.monotonic() - last_launch >= 1.0 + ): + start_host() + last_launch = time.monotonic() + time.sleep(0.05) + raise RuntimeError(f"AI Host did not become ready; log: {host_log_path()}") + + +def submit_host_job(request): + connection = ensure_host() + payload = dict(request.payload) + payload["client_id"] = _client_id + _host_http_request( + connection, + "POST", + "/jobs", + payload=payload, + timeout=5.0, + ) + return connection + + +def stop_host(): + global _host_auto_start, _host_connection, _host_snapshot + global _registered_instance_id + + connection = _host_connection + if connection is None: + connection, _health = _discover_live_host(timeout=0.5) + if connection is None: + _close_launched_host(terminate=True) + _host_auto_start = False + _host_snapshot = {"state": "STOPPED", "message": "Stopped"} + return + _register_host_client(connection) + _host_http_request( + connection, + "POST", + "/shutdown", + payload={"client_id": _client_id}, + timeout=2.0, + ) + _host_auto_start = False + deadline = time.monotonic() + 10.0 + while time.monotonic() < deadline: + try: + _check_host(connection, timeout=0.2) + except RuntimeError: + break + time.sleep(0.05) + else: + raise RuntimeError("AI Host did not stop") + _host_connection = None + _registered_instance_id = None + _close_finished_launched_host() + _host_snapshot = {"state": "STOPPED", "message": "Stopped"} + + +def restart_host(): + global _host_auto_start, _host_connection, _registered_instance_id + global _host_snapshot + + connection = ensure_host() + _host_http_request( + connection, + "POST", + "/shutdown", + payload={"client_id": _client_id}, + timeout=2.0, + ) + _host_auto_start = True + _host_connection = None + _registered_instance_id = None + _host_snapshot = {"state": "STARTING", "message": "Restarting"} + + +def enable_host(): + global _host_auto_start, _host_connection, _host_snapshot + + _host_auto_start = True + connection, health = _discover_live_host(timeout=0.2) + if connection is not None: + _host_connection = connection + _register_host_client(connection) + _host_snapshot = _format_host_snapshot(connection, health) + _redraw_host_ui() + return + start_host() + _host_snapshot = {"state": "STARTING", "message": "Starting"} + _redraw_host_ui() + + +def clear_host_models(): + connection = ensure_host() + return _host_http_request( + connection, + "POST", + "/models/clear", + timeout=5.0, + ) + + +def cancel_host_job(job_id): + connection = ensure_host() + return _host_http_request( + connection, + "DELETE", + f"/jobs/{job_id}", + payload={"client_id": _client_id}, + timeout=2.0, + ) + + +def detach_host_client(): + global _host_connection, _registered_instance_id + + connection = _host_connection + if connection is not None and _registered_instance_id is not None: + try: + _host_http_request( + connection, + "DELETE", + f"/clients/{_client_id}", + timeout=0.5, + ) + except RuntimeError: + pass + _host_connection = None + _registered_instance_id = None + _close_launched_host(terminate=False) + + +def schedule_host_start(): + timers = getattr(bpy.app, "timers", None) + if timers is None: + return + if timers.is_registered(_service_host_timer): + return + timers.register(_service_host_timer, first_interval=1.0, persistent=True) + + +def unschedule_host(): + timers = getattr(bpy.app, "timers", None) + if timers is None: + return + if timers.is_registered(_service_host_timer): + timers.unregister(_service_host_timer) + + +def host_status(): + return dict(_host_snapshot) + + +def host_busy(): + return _host_snapshot.get("state") == "BUSY" + + +def _service_host_timer(): + global _host_connection, _host_snapshot, _registered_instance_id + + if not runtime_ready(): + _host_snapshot = { + "state": "UNAVAILABLE", + "message": "Runtime is not installed", + } + return 2.0 + connection, health = _discover_live_host(timeout=0.2) + if connection is None: + _host_connection = None + _registered_instance_id = None + _close_finished_launched_host() + if not _host_auto_start: + if _host_snapshot.get("state") != "STOPPED": + _host_snapshot = {"state": "STOPPED", "message": "Stopped"} + _redraw_host_ui() + return 2.0 + try: + start_host() + except OSError as error: + _host_snapshot = {"state": "ERROR", "message": str(error)} + return 2.0 + _host_snapshot = {"state": "STARTING", "message": "Starting"} + _redraw_host_ui() + return 1.0 + try: + _host_connection = connection + if _registered_instance_id == connection.instance_id: + try: + _host_http_request( + connection, + "POST", + f"/clients/{_client_id}/heartbeat", + timeout=0.5, + ) + except RuntimeError: + _registered_instance_id = None + _register_host_client(connection) + health = _check_host(connection, timeout=0.5) + _host_snapshot = _format_host_snapshot(connection, health) + except RuntimeError as error: + _host_connection = None + _registered_instance_id = None + _host_snapshot = {"state": "ERROR", "message": str(error)} + _redraw_host_ui() + return 5.0 + + +def _register_host_client(connection): + global _registered_instance_id + + if _registered_instance_id == connection.instance_id: + return + label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip() + _host_http_request( + connection, + "POST", + "/clients", + payload={"client_id": _client_id, "label": label}, + timeout=2.0, + ) + _registered_instance_id = connection.instance_id + + +def _format_host_snapshot(connection, health): + state = "BUSY" if health.get("busy") else "READY" + return { + **health, + "state": state, + "message": "Busy" if state == "BUSY" else "Ready", + "port": connection.port, + } + + +def _redraw_host_ui(): + try: + from .status_bar import redraw_ui + + redraw_ui(bpy.context, force=True) + except (AttributeError, RuntimeError): + pass + + +def _close_finished_launched_host(): + if _launched_host is None or _launched_host.process.poll() is None: + return + _close_launched_host() + + +def _close_launched_host(terminate=False): + global _launched_host + + launched = _launched_host + _launched_host = None + if launched is None: + return + if terminate and launched.process.poll() is None: + launched.process.terminate() + try: + launched.process.wait(timeout=3.0) + except subprocess.TimeoutExpired: + launched.process.kill() + launched.process.wait(timeout=3.0) + launched.log_file.close() + + def download_model_command(model_key, status_path): model = get_downloadable_model(model_key) command = [ @@ -153,31 +604,29 @@ def download_model_command(model_key, status_path): return command -def remove_background_command( +def remove_background_request( settings, input_path, status_path, result_directory, model_key, ): - return [ - str(runtime_python()), - str(Path(__file__).with_name("worker.py")), - "remove-background", - "--input", - str(input_path), - "--model-dir", - str(model_directory(model_key)), - "--output-dir", - str(result_directory), - "--status", - str(status_path), - "--device", - settings.device.lower(), - ] + return HostJobRequest( + { + "job_id": Path(status_path).stem, + "command": "remove-background", + "parameters": { + "input": str(input_path), + "model_dir": str(model_directory(model_key)), + "output_dir": str(result_directory), + "status": str(status_path), + "device": settings.device.lower(), + }, + } + ) -def reconstruct_command( +def reconstruct_request( settings, export_format, status_path, @@ -201,53 +650,39 @@ def reconstruct_command( snapform_selected_ray_pose = ( settings.use_ray_pose if use_ray_pose is None else use_ray_pose ) - command = [ - str(runtime_python()), - str(Path(__file__).with_name("worker.py")), - "reconstruct", - "--input", - str( + parameters = { + "input": str( Path(input_path) if input_path is not None else Path(bpy.path.abspath(settings.input_path)) ), - "--model-dir", - str(model_directory(snapform_selected_key)), - "--output-dir", - str(result_directory), - "--export-format", - export_format, - "--status", - str(status_path), - "--process-res", - str(snapform_selected_process_res), - "--process-res-method", - snapform_selected_process_res_method, - "--conf-thresh-percentile", - str(settings.conf_thresh_percentile), - "--num-max-points", - str(settings.num_max_points), - "--show-cameras", - "1" if settings.show_cameras else "0", - "--gs-views-interval", - str(settings.gs_views_interval), - "--video-frames", - str(settings.video_frames), - "--device", - settings.device.lower(), - "--use-ray-pose", - "1" if snapform_selected_ray_pose else "0", - "--ref-view-strategy", - settings.ref_view_strategy, - ] + "model_dir": str(model_directory(snapform_selected_key)), + "output_dir": str(result_directory), + "export_format": export_format, + "status": str(status_path), + "process_res": snapform_selected_process_res, + "process_res_method": snapform_selected_process_res_method, + "conf_thresh_percentile": settings.conf_thresh_percentile, + "num_max_points": settings.num_max_points, + "show_cameras": 1 if settings.show_cameras else 0, + "gs_views_interval": settings.gs_views_interval, + "video_frames": settings.video_frames, + "device": settings.device.lower(), + "use_ray_pose": 1 if snapform_selected_ray_pose else 0, + "ref_view_strategy": settings.ref_view_strategy, + "remove_background_model_dir": "", + } if remove_background: - command.extend( - [ - "--remove-background-model-dir", - str(model_directory(BEN2_MODEL.key)), - ] + parameters["remove_background_model_dir"] = str( + model_directory(BEN2_MODEL.key) ) - return command + return HostJobRequest( + { + "job_id": Path(status_path).stem, + "command": "reconstruct", + "parameters": parameters, + } + ) def creation_flags(): diff --git a/src/snapform/worker.py b/src/snapform/worker.py index 49f202d..f18748f 100644 --- a/src/snapform/worker.py +++ b/src/snapform/worker.py @@ -23,10 +23,27 @@ SCENE_FILENAME = "scene.glb" EXPORT_FORMATS = ("glb", "gs_ply", "depth_vis", "mini_npz") -def write_status(path, progress, message, error=None): +class TaskCancelled(RuntimeError): + pass + + +def check_cancelled(cancel_check): + if cancel_check is not None and cancel_check(): + raise TaskCancelled("Task cancelled") + + +def write_status(path, progress, message, error=None, state=None): payload = {"progress": progress, "message": message} if error: payload["error"] = error + if state is None: + if error: + state = "failed" + elif progress >= 1.0: + state = "succeeded" + else: + state = "running" + payload["state"] = state Path(path).write_text( json.dumps(payload, ensure_ascii=False), encoding="utf-8", @@ -42,7 +59,13 @@ def collect_images(directory): ) -def extract_video_frames(video_path, output_directory, maximum_frames, cv2): +def extract_video_frames( + video_path, + output_directory, + maximum_frames, + cv2, + cancel_check=None, +): capture = cv2.VideoCapture(str(video_path)) if not capture.isOpened(): raise RuntimeError(f"Unable to open video: {video_path}") @@ -66,6 +89,7 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2): image_paths = [] try: for output_index, frame_index in enumerate(frame_indices): + check_cancelled(cancel_check) capture.set(cv2.CAP_PROP_POS_FRAMES, int(frame_index)) success, frame = capture.read() if not success: @@ -82,7 +106,8 @@ def extract_video_frames(video_path, output_directory, maximum_frames, cv2): return image_paths -def prepare_input(input_path, video_frames, temporary_root): +def prepare_input(input_path, video_frames, temporary_root, cancel_check=None): + check_cancelled(cancel_check) input_path = Path(input_path) suffix = input_path.suffix.lower() if suffix in IMAGE_SUFFIXES: @@ -102,6 +127,7 @@ def prepare_input(input_path, video_frames, temporary_root): temporary_directory, video_frames, cv2, + cancel_check=cancel_check, ) except Exception: shutil.rmtree(temporary_directory, ignore_errors=True) @@ -203,19 +229,33 @@ def select_device(requested, torch): return torch.device("cuda" if torch.cuda.is_available() else "cpu") -def infer_foreground(input_path, model_dir, device, torch): - from ben2 import BEN_Base +def infer_foreground( + input_path, + model_dir, + device, + torch, + model_cache=None, + cancel_check=None, +): from PIL import Image input_path = Path(input_path) + check_cancelled(cancel_check) if input_path.suffix.lower() not in IMAGE_SUFFIXES: raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image") - load_started = time.perf_counter() - model = BEN_Base.from_pretrained(model_dir) - model = model.to(device=device) - model.eval() - load_ms = (time.perf_counter() - load_started) * 1000.0 + owns_model = model_cache is None + if owns_model: + from ben2 import BEN_Base + load_started = time.perf_counter() + model = BEN_Base.from_pretrained(model_dir) + model = model.to(device=device) + model.eval() + load_ms = (time.perf_counter() - load_started) * 1000.0 + else: + model, load_ms = model_cache.get_ben2(model_dir, device) + + check_cancelled(cancel_check) with Image.open(input_path) as opened_image: image = opened_image.convert("RGB") if device.type == "cuda": @@ -229,17 +269,23 @@ def infer_foreground(input_path, model_dir, device, torch): if device.type == "cuda": torch.cuda.synchronize() inference_ms = (time.perf_counter() - inference_started) * 1000.0 - del model - if device.type == "cuda": - torch.cuda.empty_cache() + if owns_model: + del model + if device.type == "cuda": + torch.cuda.empty_cache() + check_cancelled(cancel_check) return foreground, load_ms, inference_ms -def remove_background(args): +def remove_background(args, model_cache=None, cancel_check=None): total_started = time.perf_counter() write_status(args.status, 0.02, "Starting BEN2 Runtime") + check_cancelled(cancel_check) - import torch + if model_cache is None: + import torch + else: + torch = model_cache.torch output_directory = Path(args.output_dir) output_directory.mkdir(parents=True, exist_ok=False) @@ -252,8 +298,11 @@ def remove_background(args): args.model_dir, device, torch, + model_cache=model_cache, + cancel_check=cancel_check, ) + check_cancelled(cancel_check) write_status(args.status, 0.92, "Writing transparent PNG") output_path = output_directory / "foreground.png" foreground.convert("RGBA").save( @@ -264,6 +313,7 @@ def remove_background(args): if not output_path.is_file(): raise RuntimeError("BEN2 did not write a foreground image") + check_cancelled(cancel_check) total_ms = (time.perf_counter() - total_started) * 1000.0 message = ( f"Background removed in {inference_ms:.0f} ms " @@ -288,10 +338,16 @@ def exported_paths(output_directory, export_format): return [] -def wait_for_exported_paths(output_directory, export_format, timeout=30.0): +def wait_for_exported_paths( + output_directory, + export_format, + timeout=30.0, + cancel_check=None, +): # 上游 NPZ exporter 使用后台线程，因此需要等待文件实际落盘。 deadline = time.monotonic() + timeout while time.monotonic() < deadline: + check_cancelled(cancel_check) result_paths = exported_paths(output_directory, export_format) if result_paths: return result_paths @@ -299,11 +355,15 @@ def wait_for_exported_paths(output_directory, export_format, timeout=30.0): return [] -def reconstruct(args): +def reconstruct(args, model_cache=None, cancel_check=None): write_status(args.status, 0.02, "Starting AI runtime") + check_cancelled(cancel_check) - import torch - from depth_anything_3.api import DepthAnything3 + if model_cache is None: + import torch + from depth_anything_3.api import DepthAnything3 + else: + torch = model_cache.torch from depth_anything_3.utils.export import export output_directory = Path(args.output_dir) @@ -313,9 +373,11 @@ def reconstruct(args): args.input, args.video_frames, output_directory.parent, + cancel_check=cancel_check, ) try: + check_cancelled(cancel_check) device = select_device(args.device, torch) remove_background_model_dir = getattr( args, @@ -329,7 +391,10 @@ def reconstruct(args): remove_background_model_dir, device, torch, + model_cache=model_cache, + cancel_check=cancel_check, ) + check_cancelled(cancel_check) foreground.convert("RGBA").save( output_directory / "foreground.png", format="PNG", @@ -342,10 +407,14 @@ def reconstruct(args): f"model load {load_ms:.0f} ms" ) write_status(args.status, 0.15, f"Loading model on {device.type}") - model = DepthAnything3.from_pretrained(args.model_dir) - model = model.to(device=device) - model.eval() + if model_cache is None: + model = DepthAnything3.from_pretrained(args.model_dir) + model = model.to(device=device) + model.eval() + else: + model = model_cache.get_da3(args.model_dir, device) + check_cancelled(cancel_check) write_status( args.status, 0.25, @@ -372,6 +441,7 @@ def reconstruct(args): use_ray_pose=bool(args.use_ray_pose), ref_view_strategy=args.ref_view_strategy, ) + check_cancelled(cancel_check) write_status(args.status, 0.88, f"Exporting {args.export_format}") export( prediction, @@ -379,15 +449,18 @@ def reconstruct(args): str(output_directory), **export_kwargs, ) + check_cancelled(cancel_check) result_paths = wait_for_exported_paths( output_directory, args.export_format, + cancel_check=cancel_check, ) if not result_paths: raise RuntimeError( f"No {args.export_format} files were exported" ) + check_cancelled(cancel_check) write_status( args.status, 1.0, diff --git a/tests/test_addon.py b/tests/test_addon.py index 2135985..69f805d 100644 --- a/tests/test_addon.py +++ b/tests/test_addon.py @@ -343,7 +343,7 @@ class AddonTest(unittest.TestCase): self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"]) self.assertEqual( operator.__annotations__["snapform_process_res_factor"]["default"], - 0.5, + 1.0, ) self.assertTrue( operator.__annotations__["snapform_delete_original"]["default"] @@ -403,7 +403,7 @@ class AddonTest(unittest.TestCase): '"snapform_process_res_factor",' ), draw_source.index( - 'layout.prop(self, "snapform_subdivision_level")' + 'settings_body.prop(self, "snapform_subdivision_level")' ), ) self.assertIn( @@ -420,17 +420,30 @@ class AddonTest(unittest.TestCase): draw_source, ) self.assertIn( - 'layout.prop(self, "snapform_remove_background")', + 'default_closed=True', draw_source, ) + self.assertIn('settings_header.label(text="Options")', draw_source) + self.assertIn('settings_header.label(text="Options")', draw_source) self.assertNotIn("layout.separator", draw_source) - self.assertNotIn("layout.label", draw_source) - self.assertIn( - 'layout.prop(self, "snapform_unshaded")', - draw_source, + setting_fields = ( + "snapform_displacement_via", + "snapform_subdivision_level", + "snapform_remove_background", + "snapform_unshaded", ) - model_items = operator.__annotations__["snapform_model"]["items"] - model_items = model_items(None, None) + field_positions = [ + draw_source.index(f'settings_body.prop(self, "{field}")') + for field in setting_fields + ] + self.assertEqual(field_positions, sorted(field_positions)) + model_items_callback = operator.__annotations__["snapform_model"]["items"] + with patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ): + model_items = model_items_callback(None, None) self.assertEqual( [item[0] for item in model_items], [ @@ -459,11 +472,33 @@ class AddonTest(unittest.TestCase): ) self.addon.image_depth._last_depth_plane_model = "GIANT_1_1" - result = operator.invoke(context, None) + with patch.object( + self.addon.runtime, + "model_ready", + side_effect=lambda key: key == "GIANT_1_1", + ): + result = operator.invoke(context, None) self.assertIs(result, operator) self.assertEqual(operator.snapform_model, "GIANT_1_1") + def test_conversion_dialog_only_lists_installed_models(self): + items_callback = self.addon.ConvertImageToDepthPlane.__annotations__[ + "snapform_model" + ]["items"] + with patch.object( + self.addon.runtime, + "model_ready", + side_effect=lambda key: key in {"SMALL", "BASE"}, + ): + items = items_callback(None, None) + + self.assertEqual([item[0] for item in items], ["SMALL", "BASE"]) + draw_source = inspect.getsource(self.addon.ConvertImageToDepthPlane.draw) + self.assertIn('text="No DA3 models installed"', draw_source) + self.assertIn('text="Download one in Host > Models"', draw_source) + self.assertNotIn("_dialog_requirements_ready", inspect.getsource(self.addon.ConvertImageToDepthPlane)) + def test_depth_conversion_progress_is_provider_neutral(self): starting_source = inspect.getsource( self.addon.RunJob._starting_message @@ -532,7 +567,11 @@ class AddonTest(unittest.TestCase): self.addon.image_depth, "save_packed_image", return_value=packed_path, - ) as save_packed_image: + ) as save_packed_image, patch.object( + self.addon.runtime, + "model_ready", + return_value=True, + ): result = operator.execute(context) self.assertEqual(result, {"FINISHED"}) @@ -640,8 +679,10 @@ class AddonTest(unittest.TestCase): return_value=True, ) as model_ready, patch.object( self.addon.runtime, - "reconstruct_command", - return_value=["worker"], + "reconstruct_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), ) as reconstruct: operator.snapform_model_key = snapform_model operator.snapform_use_ray_pose = snapform_ray_pose @@ -650,7 +691,10 @@ class AddonTest(unittest.TestCase): settings, ) - self.assertEqual(command, ["worker"]) + self.assertIsInstance( + command, + self.addon.runtime.HostJobRequest, + ) self.assertEqual(prepared_paths, paths) model_ready.assert_called_once_with(snapform_model) self.assertEqual(reconstruct.call_args.args[1], "mini_npz") @@ -713,7 +757,7 @@ class AddonTest(unittest.TestCase): with self.assertRaisesRegex(RuntimeError, "BEN2 Base"): operator._prepare_command(None, settings) - def test_image_depth_background_option_reaches_worker_command(self): + def test_image_depth_background_option_reaches_host_request(self): runtime = importlib.import_module("snapform.runtime") settings = SimpleNamespace( input_path="input.png", @@ -739,7 +783,7 @@ class AddonTest(unittest.TestCase): "model_directory", side_effect=lambda key: Path(key.lower()), ): - command = runtime.reconstruct_command( + request = runtime.reconstruct_request( settings, "mini_npz", Path("status.json"), @@ -747,8 +791,10 @@ class AddonTest(unittest.TestCase): remove_background=True, ) - option_index = command.index("--remove-background-model-dir") - self.assertEqual(command[option_index + 1], str(Path("ben2_base"))) + self.assertEqual( + request.payload["parameters"]["remove_background_model_dir"], + str(Path("ben2_base")), + ) def test_preferences_use_package_identifier(self): self.assertEqual( @@ -767,6 +813,25 @@ class AddonTest(unittest.TestCase): str(runtime.DEFAULT_STORAGE_ROOT), ) + def test_ai_host_uses_authenticated_loopback_http(self): + runtime = importlib.import_module("snapform.runtime") + request_source = inspect.getsource(runtime._host_http_request) + start_source = inspect.getsource(runtime.start_host) + + self.assertEqual(runtime.HOST_ADDRESS, "127.0.0.1") + self.assertIn('"X-SnapForm-Token"', request_source) + self.assertIn('"/health"', inspect.getsource(runtime._check_host)) + self.assertIn('with_name("ai_host.py")', start_source) + self.assertIn("ProxyHandler({})", inspect.getsource(runtime)) + self.assertIn('return storage_root() / "host.json"', inspect.getsource(runtime.host_discovery_path)) + + def test_host_jobs_complete_from_status_state(self): + source = inspect.getsource(self.addon.RunJob.modal) + + self.assertIn("if job.uses_host:", source) + self.assertIn('state == "succeeded"', source) + self.assertIn('{"failed", "cancelled"}', source) + def test_process_res_method_defaults_to_upper_bound_resize(self): process_res = self.addon.SnapFormSettings.__annotations__[ "process_res" @@ -780,7 +845,7 @@ class AddonTest(unittest.TestCase): self.assertEqual(process_res["default"], 504) self.assertEqual(process_res_method["default"], "upper_bound_resize") - def test_reconstruct_command_preserves_da3_parameter_names(self): + def test_reconstruct_request_preserves_da3_parameter_names(self): runtime = importlib.import_module("snapform.runtime") settings = SimpleNamespace( input_path="input.png", @@ -806,29 +871,30 @@ class AddonTest(unittest.TestCase): "model_directory", return_value=Path("model"), ): - command = runtime.reconstruct_command( + request = runtime.reconstruct_request( settings, "gs_ply", Path("status.json"), Path("result"), ) - self.assertIn("--export-format", command) - self.assertIn("gs_ply", command) + parameters = request.payload["parameters"] + self.assertEqual(request.payload["command"], "reconstruct") + self.assertEqual(parameters["export_format"], "gs_ply") for parameter in ( - "--process-res", - "--process-res-method", - "--conf-thresh-percentile", - "--num-max-points", - "--show-cameras", - "--gs-views-interval", - "--use-ray-pose", - "--ref-view-strategy", + "process_res", + "process_res_method", + "conf_thresh_percentile", + "num_max_points", + "show_cameras", + "gs_views_interval", + "use_ray_pose", + "ref_view_strategy", ): with self.subTest(parameter=parameter): - self.assertIn(parameter, command) + self.assertIn(parameter, parameters) - def test_reconstruct_command_accepts_image_depth_overrides(self): + def test_reconstruct_request_accepts_image_depth_overrides(self): runtime = importlib.import_module("snapform.runtime") settings = SimpleNamespace( input_path="input.png", @@ -854,7 +920,7 @@ class AddonTest(unittest.TestCase): "model_directory", return_value=Path("large-model"), ) as snapform_model_directory: - snapform_command = runtime.reconstruct_command( + request = runtime.reconstruct_request( settings, "mini_npz", Path("status.json"), @@ -866,15 +932,13 @@ class AddonTest(unittest.TestCase): ) snapform_model_directory.assert_called_once_with("BASE") - snapform_process_res_index = snapform_command.index("--process-res") - self.assertEqual(snapform_command[snapform_process_res_index + 1], "252") - snapform_process_res_method_index = snapform_command.index("--process-res-method") + parameters = request.payload["parameters"] + self.assertEqual(parameters["process_res"], 252) self.assertEqual( - snapform_command[snapform_process_res_method_index + 1], + parameters["process_res_method"], "lower_bound_resize", ) - snapform_ray_pose_index = snapform_command.index("--use-ray-pose") - self.assertEqual(snapform_command[snapform_ray_pose_index + 1], "0") + self.assertEqual(parameters["use_ray_pose"], 0) def test_export_format_property_includes_all_da3_formats(self): export_format = self.addon.SnapFormSettings.__annotations__[ @@ -904,37 +968,139 @@ class AddonTest(unittest.TestCase): "DEFAULT_CLOSED", self.addon.DebugPanel.bl_options, ) + for panel in (self.addon.Da3DebugPanel, self.addon.Ben2DebugPanel): + self.assertEqual(panel.bl_parent_id, self.addon.DebugPanel.bl_idname) + self.assertIn("DEFAULT_CLOSED", panel.bl_options) def test_debug_panel_uses_one_export_button_and_dynamic_settings(self): - draw_source = inspect.getsource( - self.addon.DebugPanel.draw - ) + debug_source = inspect.getsource(self.addon.DebugPanel.draw) + da3_source = inspect.getsource(self.addon.Da3DebugPanel.draw) + ben2_source = inspect.getsource(self.addon.Ben2DebugPanel.draw) - self.assertEqual(draw_source.count("export_row.operator("), 1) - self.assertIn('settings.export_format == "glb"', draw_source) - self.assertIn('settings.export_format == "gs_ply"', draw_source) - self.assertIn('"gs_model" if settings.export_format', draw_source) - self.assertIn("layout.prop(settings, model_property)", draw_source) - self.assertNotIn("model_row", draw_source) - self.assertIn("export_row.enabled = active_job() is None", draw_source) - self.assertNotIn('text="Ready"', draw_source) - self.assertNotIn('text="Missing"', draw_source) - self.assertNotIn("License:", draw_source) - self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source) - self.assertIn('text="BEN2 Background Removal"', draw_source) + self.assertIn('layout.prop(settings, "device")', debug_source) + self.assertEqual(da3_source.count("export_row.operator("), 1) + self.assertIn('settings.export_format == "glb"', da3_source) + self.assertIn('settings.export_format == "gs_ply"', da3_source) + self.assertIn('"gs_model" if settings.export_format', da3_source) + self.assertIn("layout.prop(settings, model_property)", da3_source) + self.assertNotIn("model_row", da3_source) + self.assertIn("and model_ready(selected_model_key(settings))", da3_source) + self.assertNotIn('text="Ready"', da3_source) + self.assertNotIn('text="Missing"', da3_source) + self.assertNotIn("License:", da3_source) + self.assertIn('operator.job_type = "RECONSTRUCT"', da3_source) self.assertIn( - 'background_box.prop(settings, "background_input_path")', - draw_source, + 'layout.prop(settings, "background_input_path")', + ben2_source, ) - self.assertIn('text="Remove Background"', draw_source) - self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', draw_source) + self.assertIn('text="Remove Background"', ben2_source) + self.assertIn("and model_ready(BEN2_MODEL.key)", ben2_source) + self.assertIn('operator.job_type = "REMOVE_BACKGROUND"', ben2_source) - def test_runtime_models_panel_assigns_each_download_target(self): + def test_host_panel_exposes_shared_host_state_and_controls(self): + self.assertIn(self.addon.HostPanel, self.addon.CLASSES) + self.assertIn(self.addon.ControlHost, self.addon.CLASSES) + source = inspect.getsource(self.addon.HostPanel.draw) + + self.assertEqual(self.addon.HostPanel.bl_label, "Host") + self.assertIn('text="Runtime is not installed"', source) + self.assertIn('text="Install Runtime"', source) + self.assertIn("runtime.host_status()", source) + self.assertIn('f"Running (Clients:', source) + self.assertIn('else "Running (Busy)"', source) + self.assertIn('"READY": "RADIOBUT_ON"', source) + self.assertIn('"BUSY": "TIME"', source) + self.assertNotIn('"READY": "PLAY"', source) + self.assertNotIn('text=f"Task:', source) + self.assertIn('elif state == "STOPPED":', source) + self.assertNotIn("Blender Clients:", source) + self.assertNotIn("status['port']", source) + self.assertIn('icon="TRASH"', source) + self.assertIn('icon="FILE_REFRESH"', source) + self.assertIn('icon="PAUSE"', source) + self.assertNotIn('text="Open Log"', source) + preferences_source = inspect.getsource(self.addon.SnapFormPreferences.draw) + self.assertIn('text="Open Host Log"', preferences_source) + + def test_host_start_control_does_not_wait_for_model_runtime(self): + source = inspect.getsource(self.addon.runtime.enable_host) + + self.assertIn("start_host()", source) + self.assertNotIn("ensure_host()", source) + + def test_ai_operation_restarts_a_manually_stopped_host(self): + runtime = self.addon.runtime + connection = runtime.HostConnection(51234, "token", 42, "host") + health = { + "host_schema": runtime.HOST_SCHEMA, + "instance_id": "host", + "busy": False, + "clients": 1, + "cached_models": {}, + } + previous_auto_start = runtime._host_auto_start + previous_connection = runtime._host_connection + previous_snapshot = runtime._host_snapshot + runtime._host_auto_start = False + runtime._host_connection = None + runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"} + try: + with patch.object( + runtime, + "_discover_live_host", + return_value=(connection, health), + ), patch.object( + runtime, + "_register_host_client", + ), patch.object( + runtime, + "_check_host", + return_value=health, + ), patch.object(runtime, "_redraw_host_ui"): + result = runtime.ensure_host(timeout=0.1) + + self.assertIs(result, connection) + self.assertTrue(runtime._host_auto_start) + self.assertEqual(runtime.host_status()["state"], "READY") + finally: + runtime._host_auto_start = previous_auto_start + runtime._host_connection = previous_connection + runtime._host_snapshot = previous_snapshot + + def test_blender_startup_does_not_start_ai_host(self): + runtime = self.addon.runtime + previous_auto_start = runtime._host_auto_start + previous_snapshot = runtime._host_snapshot + runtime._host_auto_start = False + runtime._host_snapshot = {"state": "STOPPED", "message": "Stopped"} + try: + with patch.object( + runtime, + "runtime_ready", + return_value=True, + ), patch.object( + runtime, + "_discover_live_host", + return_value=(None, None), + ), patch.object( + runtime, + "_close_finished_launched_host", + ), patch.object(runtime, "start_host") as start_host: + interval = runtime._service_host_timer() + + start_host.assert_not_called() + self.assertEqual(interval, 2.0) + self.assertEqual(runtime.host_status()["state"], "STOPPED") + finally: + runtime._host_auto_start = previous_auto_start + runtime._host_snapshot = previous_snapshot + + def test_models_panel_is_collapsed_and_assigns_each_download_target(self): draw_source = inspect.getsource( - self.addon.RuntimeModelsPanel.draw + self.addon.ModelsPanel.draw ) model_source = inspect.getsource( - self.addon.RuntimeModelsPanel._draw_model + self.addon.ModelsPanel._draw_model ) self.assertIn("for model in MODELS.values()", draw_source) @@ -951,8 +1117,16 @@ class AddonTest(unittest.TestCase): self.assertNotIn("Isolated Runtime ready", draw_source) self.assertNotIn("License:", model_source) self.assertEqual( - self.addon.RuntimeModelsPanel.bl_label, - "Setup", + self.addon.ModelsPanel.bl_label, + "Models", + ) + self.assertEqual( + self.addon.ModelsPanel.bl_parent_id, + self.addon.HostPanel.bl_idname, + ) + self.assertEqual( + self.addon.ModelsPanel.bl_options, + {"DEFAULT_CLOSED"}, ) def test_ben2_debug_job_uses_its_image_and_model(self): @@ -987,15 +1161,17 @@ class AddonTest(unittest.TestCase): return_value=True, ) as model_ready, patch.object( self.addon.runtime, - "remove_background_command", - return_value=["worker"], + "remove_background_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), ) as remove_background: command, prepared_paths = operator._prepare_command( None, settings, ) - self.assertEqual(command, ["worker"]) + self.assertIsInstance(command, self.addon.runtime.HostJobRequest) self.assertEqual(prepared_paths, paths) model_ready.assert_called_once_with("BEN2_BASE") self.assertEqual(remove_background.call_args.args[1], input_path) @@ -1030,15 +1206,17 @@ class AddonTest(unittest.TestCase): return_value=True, ), patch.object( self.addon.runtime, - "remove_background_command", - return_value=["worker"], + "remove_background_request", + return_value=self.addon.runtime.HostJobRequest( + {"job_id": "test"} + ), ) as remove_background: command, prepared_paths = operator._prepare_command( None, settings, ) - self.assertEqual(command, ["worker"]) + self.assertIsInstance(command, self.addon.runtime.HostJobRequest) self.assertEqual(prepared_paths, paths) self.assertEqual(remove_background.call_args.args[1], input_path) diff --git a/tests/test_ai_host.py b/tests/test_ai_host.py new file mode 100644 index 0000000..e74669a --- /dev/null +++ b/tests/test_ai_host.py @@ -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() diff --git a/tests/test_bootstrap_runtime.py b/tests/test_bootstrap_runtime.py index f390541..1949d78 100644 --- a/tests/test_bootstrap_runtime.py +++ b/tests/test_bootstrap_runtime.py @@ -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__": diff --git a/tests/test_pack.py b/tests/test_pack.py index 7f3b079..69dca31 100644 --- a/tests/test_pack.py +++ b/tests/test_pack.py @@ -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) diff --git a/tests/test_worker.py b/tests/test_worker.py index 1aaa7e9..8770072 100644 --- a/tests/test_worker.py +++ b/tests/test_worker.py @@ -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)