Enhance device management and model handling in SnapForm
- Updated frame extraction logic to allow for single frame extraction. - Integrated model readiness checks for BEN2 and DA3 models in relevant operators. - Added directory name attribute to model specifications for better organization. - Improved device selection logic for production and debug tasks. - Refactored preferences to expose device settings and ensure they reflect available hardware. - Enhanced runtime device detection and management, including support for CUDA and MPS. - Updated tests to cover new device management features and model directory handling. - Removed deprecated download model tests to streamline codebase.
This commit is contained in:
@@ -18,8 +18,8 @@ SnapForm 是一个 Blender 扩展,可以把一张照片、一段视频或一
|
|||||||
- Linux:`SnapForm.v0.10.10.linux-x64.zip`
|
- Linux:`SnapForm.v0.10.10.linux-x64.zip`
|
||||||
2. 打开 Blender,进入 `Preferences > Extensions`
|
2. 打开 Blender,进入 `Preferences > Extensions`
|
||||||
3. 点击 `Install from Disk`,选择下载的压缩包
|
3. 点击 `Install from Disk`,选择下载的压缩包
|
||||||
4. 在扩展偏好设置中确认 `uv` 路径与 Torch 后端(保持默认即可);如果下载经常超时,
|
4. 在扩展偏好设置中确认 Storage Root 与 Device(保持默认即可);Device 默认取本机
|
||||||
可以在 `PyPI Index` 里填一个国内镜像,例如 `https://pypi.tuna.tsinghua.edu.cn/simple`
|
最快的 CUDA/MPS,否则 CPU。下载超时会自动重试国内镜像,无需手动配置
|
||||||
5. 打开 3D 视图右侧的 `SnapForm` 面板,点击 `Install Runtime`
|
5. 打开 3D 视图右侧的 `SnapForm` 面板,点击 `Install Runtime`
|
||||||
|
|
||||||
DA3、BEN2 源码包和 uv 已内置在扩展里,安装时直接从本地解包,不需要访问 GitHub。首次使用仍会自动
|
DA3、BEN2 源码包和 uv 已内置在扩展里,安装时直接从本地解包,不需要访问 GitHub。首次使用仍会自动
|
||||||
|
|||||||
@@ -50,18 +50,21 @@
|
|||||||
`linux-x64` 四个平台的 Extension ZIP(如 `SnapForm.v0.10.10.macos-arm64.zip`)。
|
`linux-x64` 四个平台的 Extension ZIP(如 `SnapForm.v0.10.10.macos-arm64.zip`)。
|
||||||
每个包只含对应平台的 uv 二进制,并内置 DA3、BEN2 源码包。
|
每个包只含对应平台的 uv 二进制,并内置 DA3、BEN2 源码包。
|
||||||
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
|
2. 在 Blender 的 `Preferences > Extensions` 中选择 `Install from Disk`。
|
||||||
3. 在 Extension Preferences 中确认 `uv` 路径和 Torch Backend。
|
3. 在 Extension Preferences 中确认 Storage Root 与 Device。Device 是正式功能
|
||||||
|
(Image Empty 右键菜单)使用的推理设备,默认取本机最快的 CUDA/MPS,否则 CPU,且只列出
|
||||||
|
本机支持的设备。
|
||||||
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn;
|
4. 在 `Host` 点击 `Install Runtime`。首次安装会下载 Python、PyTorch、FastAPI、Uvicorn;
|
||||||
DA3、BEN2 源码包与 uv 已内置在扩展中,安装时从 `~/.snapform/sources/` 解包使用,
|
DA3、BEN2 源码包与 uv 已内置在扩展中,安装时从 `~/.snapform/sources/` 解包使用,
|
||||||
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度
|
并默认下载 `DA3 Small`、`DA3 Base` 和 `BEN2 Base` 模型。安装按 8 个阶段依次推进,每段进度
|
||||||
从 0 到 100%:解包内置 uv、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过,
|
从 0 到 100%:解包内置 uv、创建 Python Runtime、下载 PyTorch、安装 xFormers(macOS 跳过,
|
||||||
内置 DA3 源码包为官方原版并已移除 xformers 硬依赖,自动回退纯 PyTorch 实现)、安装 AI 依赖,
|
内置 DA3 源码包为官方原版并已移除 xformers 硬依赖,自动回退纯 PyTorch 实现)、安装 AI 依赖,
|
||||||
然后是三个模型的下载。进度文字统一为 `阶段/总数 阶段名: 状态`,
|
然后是三个模型的下载。安装完成后进入 `验证 Runtime` 阶段,会导入 Torch 与全部 AI 依赖;
|
||||||
|
macOS 首次导入 Torch 会初始化 Metal 后端(编译 shader 缓存),该阶段可能持续较久,状态栏
|
||||||
|
会显示已用时间而不是卡住。进度文字统一为 `阶段/总数 阶段名: 状态`,
|
||||||
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。
|
模型下载阶段会实时显示下载速度,例如 `6/8 da3 small: downloading 14.2MB/s`。
|
||||||
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。
|
Runtime 就绪后即使模型下载中断也不影响 Host 使用,可稍后在 `Models` 面板单独补下。
|
||||||
安装不做启动网络预检:若网络无法稳定访问 pypi.org,可在扩展偏好设置的 `PyPI Index`
|
安装不做启动网络预检:pip 安装失败会自动依次重试清华、阿里镜像,模型下载由 worker
|
||||||
中指定镜像源(如清华、阿里),该地址会通过 `--pypi-index` 传给 bootstrap 并优先使用;
|
自行探测回退到 HF 镜像;需要手动指定镜像时可通过 `PIP_INDEX_URL` 环境变量覆盖。
|
||||||
未指定时,pip 安装失败会自动重试一次清华镜像,模型下载由 worker 自行探测回退到 HF 镜像。
|
|
||||||
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
|
5. 在 3D View 按 `N`,打开 `SnapForm` 页签。
|
||||||
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
|
6. 如需其他模型,展开 `Host > Models`,找到模型并点击该行的 `Download`。
|
||||||
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
|
7. 展开默认收起的 `Debug`,选择 `Export Format` 和本次执行使用的模型。
|
||||||
@@ -85,7 +88,7 @@
|
|||||||
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
|
状态和原始内容。BEN2 的 `foreground.png` 也会保留在对应 job 目录中。
|
||||||
|
|
||||||
Image Empty 引用视频或序列帧时,Remove Background 直接执行,默认取 Empty 的 Image Frames
|
Image Empty 引用视频或序列帧时,Remove Background 直接执行,默认取 Empty 的 Image Frames
|
||||||
数值(未设置时为完整帧数),起始帧跟随 Empty 的 Offset;若 Offset 加 Frames 超出视频长度,
|
数值(未设置时为完整帧数,显式设为 1 时只处理 1 帧),起始帧跟随 Empty 的 Offset;若 Offset 加 Frames 超出视频长度,
|
||||||
会把窗口整体收进视频内(保证实际处理 Frames 帧数),Offset 过小时从第 1 帧开始。输出
|
会把窗口整体收进视频内(保证实际处理 Frames 帧数),Offset 过小时从第 1 帧开始。输出
|
||||||
`foreground.0001.png` 等透明序列帧(帧号从 0001 开始,与 Blender 场景帧一一对应),并把
|
`foreground.0001.png` 等透明序列帧(帧号从 0001 开始,与 Blender 场景帧一一对应),并把
|
||||||
源帧以 `frame.0001.jpg` 等命名复制到 job 目录根下供材质颜色导入。替换 Empty 的 Image 为
|
源帧以 `frame.0001.jpg` 等命名复制到 job 目录根下供材质颜色导入。替换 Empty 的 Image 为
|
||||||
@@ -115,7 +118,7 @@ Cyclic 沿用原 Empty 的设置。序列帧无法被 Blender 整体 Pack,因
|
|||||||
`img.foreground` 与 `img.depth`,不会生成 `img.foreground.depth`。Original 仍用原图作为
|
`img.foreground` 与 `img.depth`,不会生成 `img.foreground.depth`。Original 仍用原图作为
|
||||||
位移来源,不改变原有深度逻辑。
|
位移来源,不改变原有深度逻辑。
|
||||||
Image Empty 引用视频或序列帧时,AI Generated 会按 Empty 的 Image Frames 数值逐帧估计
|
Image Empty 引用视频或序列帧时,AI Generated 会按 Empty 的 Image Frames 数值逐帧估计
|
||||||
深度(未设置时为完整帧数),起始帧跟随 Empty 的 Offset,与 Remove Background 范围一致;
|
深度(未设置时为完整帧数,显式设为 1 时只处理 1 帧),起始帧跟随 Empty 的 Offset,与 Remove Background 范围一致;
|
||||||
Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty
|
Offset 超出视频长度时同样把窗口收进视频内。材质节点与置换贴图的 Start 保持与原 Empty
|
||||||
一致;Original 模式(无 Remove Background)时,单帧会复制原图改名为 `.depth` 供位移使用,
|
一致;Original 模式(无 Remove Background)时,单帧会复制原图改名为 `.depth` 供位移使用,
|
||||||
视频或序列帧则抽帧保存为 job 目录根下的 `frame.0001.jpg` 等,并作为 `.depth` 序列使用。
|
视频或序列帧则抽帧保存为 job 目录根下的 `frame.0001.jpg` 等,并作为 `.depth` 序列使用。
|
||||||
@@ -307,8 +310,12 @@ Windows 上代理若只在系统设置里而未写入环境变量,安装程序
|
|||||||
并核对覆盖文件是否仍适用
|
并核对覆盖文件是否仍适用
|
||||||
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
|
- `addict` 2.4.0(补充 DA3 上游包元数据暂未声明的运行依赖)
|
||||||
|
|
||||||
默认 CUDA Backend 为 CUDA 12.8,也可以选择 CPU。CPU 可以用于验证,但实际推理
|
安装 Runtime 时会自动选择 PyTorch 后端:检测到 NVIDIA 驱动则安装 CUDA 12.8 版本,否则安装
|
||||||
通常很慢。
|
CPU 版;macOS 始终安装 CPU 版 PyTorch(MPS 包含其中)。CPU 可以用于验证,但实际推理通常很慢。
|
||||||
|
Extension Preferences 的 `Device` 决定正式功能的推理设备,默认取本机最快的 CUDA/MPS,否则
|
||||||
|
CPU;Debug 面板的 `Debug Device` 只影响 Debug 任务。两处设备选项都只列出本机支持的设备:
|
||||||
|
Windows 不显示 MPS,机器上没有可用 NVIDIA 驱动时不显示 CUDA。Runtime 安装完成时会检测真实
|
||||||
|
可用的设备并写入 `snapform_ready.json`,设备列表以该检测结果为准。
|
||||||
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
|
CUDA 12.8 的 PyTorch wheel 自带 CUDA 运行库,无需安装 CUDA Toolkit;只需要
|
||||||
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
|
NVIDIA 驱动 ≥ 570(数据中心显卡可用更老驱动)。驱动不满足或没有 NVIDIA 显卡时,
|
||||||
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
|
设备会自动回退到 CPU(macOS 为 MPS),只有显式选择 CUDA 才会报错。
|
||||||
@@ -352,6 +359,9 @@ Runtime、模型和任务结果默认保存在用户主目录的 `~/.snapform`
|
|||||||
└── host.lock # Host 运行时持有
|
└── host.lock # Host 运行时持有
|
||||||
```
|
```
|
||||||
|
|
||||||
|
`models/` 下每个模型一个文件夹,命名统一带前缀:DA3 为 `da3_*`(如 `da3_small`、`da3_base`),
|
||||||
|
BEN2 为 `ben2_*`(如 `ben2_base`)。旧版无前缀的 DA3 文件夹会在首次使用时自动改名为新名称。
|
||||||
|
|
||||||
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
|
不同 Blender 版本安装相同 Extension 后会共用这个目录。也可以在 Extension
|
||||||
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
|
Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Extension 安装目录,
|
||||||
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
|
因为该目录可能只读,并且升级时会被替换。不要同时从多个 Blender 实例执行
|
||||||
@@ -360,7 +370,7 @@ Preferences 的 `Storage Root` 中修改共享位置。不要把模型写入 Ext
|
|||||||
## 开发
|
## 开发
|
||||||
|
|
||||||
```bash
|
```bash
|
||||||
uv run python -m unittest discover -s tests -v
|
uv run pytest tests -v
|
||||||
uv run python pack.py
|
uv run python pack.py
|
||||||
```
|
```
|
||||||
|
|
||||||
|
|||||||
@@ -7,5 +7,8 @@ requires-python = ">=3.11"
|
|||||||
license = "GPL-3.0-or-later"
|
license = "GPL-3.0-or-later"
|
||||||
dependencies = []
|
dependencies = []
|
||||||
|
|
||||||
|
[dependency-groups]
|
||||||
|
dev = ["pytest"]
|
||||||
|
|
||||||
[tool.uv]
|
[tool.uv]
|
||||||
package = false
|
package = false
|
||||||
|
|||||||
@@ -147,7 +147,7 @@ class ModelCache:
|
|||||||
Path(self.ben2_key[0]).name if self.ben2_key is not None else ""
|
Path(self.ben2_key[0]).name if self.ben2_key is not None else ""
|
||||||
),
|
),
|
||||||
"da3": (
|
"da3": (
|
||||||
f"da3_{Path(self.da3_key[0]).name}"
|
Path(self.da3_key[0]).name
|
||||||
if self.da3_key is not None
|
if self.da3_key is not None
|
||||||
else ""
|
else ""
|
||||||
),
|
),
|
||||||
|
|||||||
@@ -50,11 +50,15 @@ PLAIN_SLICE_SPAN = 0.02
|
|||||||
RAMP_RATE = 0.02
|
RAMP_RATE = 0.02
|
||||||
CREEP_CAP = 0.60
|
CREEP_CAP = 0.60
|
||||||
CHECK_CODE = (
|
CHECK_CODE = (
|
||||||
"import torch; "
|
"import json; import torch; "
|
||||||
"from depth_anything_3.api import DepthAnything3; "
|
"from depth_anything_3.api import DepthAnything3; "
|
||||||
"from ben2 import BEN_Base; "
|
"from ben2 import BEN_Base; "
|
||||||
"import fastapi, uvicorn; "
|
"import fastapi, uvicorn; "
|
||||||
"print(torch.__version__, torch.cuda.is_available())"
|
"mps = getattr(torch.backends, 'mps', None); "
|
||||||
|
"print('SNAPFORM_DEVICES=' + json.dumps({"
|
||||||
|
"'torch': torch.__version__, "
|
||||||
|
"'cuda': torch.cuda.is_available(), "
|
||||||
|
"'mps': bool(mps is not None and mps.is_available())}))"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
@@ -373,7 +377,7 @@ def install_stages(args, runtime_dir, python, uv):
|
|||||||
"--repository",
|
"--repository",
|
||||||
model.repository,
|
model.repository,
|
||||||
"--model-dir",
|
"--model-dir",
|
||||||
models_dir / model.key.lower(),
|
models_dir / model.directory_name,
|
||||||
"--status",
|
"--status",
|
||||||
args.status,
|
args.status,
|
||||||
]
|
]
|
||||||
@@ -405,6 +409,70 @@ def verify_command(python):
|
|||||||
return [python, "-c", CHECK_CODE]
|
return [python, "-c", CHECK_CODE]
|
||||||
|
|
||||||
|
|
||||||
|
def run_verify(
|
||||||
|
python,
|
||||||
|
status_path,
|
||||||
|
start=VERIFY_SLICE_START,
|
||||||
|
end=VERIFY_SLICE_END,
|
||||||
|
clock=time.monotonic,
|
||||||
|
):
|
||||||
|
"""Run the runtime import check with heartbeat progress; return the device report."""
|
||||||
|
command = verify_command(python)
|
||||||
|
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
|
||||||
|
process = subprocess.Popen(
|
||||||
|
[str(value) for value in command],
|
||||||
|
stdout=subprocess.PIPE,
|
||||||
|
stderr=subprocess.STDOUT,
|
||||||
|
text=True,
|
||||||
|
encoding="utf-8",
|
||||||
|
errors="replace",
|
||||||
|
)
|
||||||
|
output_queue = queue.Queue()
|
||||||
|
|
||||||
|
def read_output():
|
||||||
|
for raw_line in process.stdout:
|
||||||
|
output_queue.put(raw_line)
|
||||||
|
output_queue.put(None)
|
||||||
|
|
||||||
|
threading.Thread(target=read_output, daemon=True).start()
|
||||||
|
|
||||||
|
report = {}
|
||||||
|
started = clock()
|
||||||
|
last_heartbeat = 0.0
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
raw_line = output_queue.get(timeout=0.2)
|
||||||
|
empty = False
|
||||||
|
except queue.Empty:
|
||||||
|
raw_line = None
|
||||||
|
empty = True
|
||||||
|
if raw_line is None and not empty:
|
||||||
|
break
|
||||||
|
if empty:
|
||||||
|
elapsed = clock() - started
|
||||||
|
if elapsed - last_heartbeat >= 5.0:
|
||||||
|
last_heartbeat = elapsed
|
||||||
|
creep = min(elapsed / 60.0, CREEP_CAP)
|
||||||
|
progress = start + creep * (end - start)
|
||||||
|
write_status(
|
||||||
|
status_path,
|
||||||
|
progress,
|
||||||
|
f"verifying runtime {format_elapsed(elapsed)}",
|
||||||
|
)
|
||||||
|
continue
|
||||||
|
line = raw_line.rstrip()
|
||||||
|
print(line)
|
||||||
|
if line.startswith("SNAPFORM_DEVICES="):
|
||||||
|
try:
|
||||||
|
report = json.loads(line[len("SNAPFORM_DEVICES=") :])
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
report = {}
|
||||||
|
return_code = process.wait()
|
||||||
|
if return_code != 0:
|
||||||
|
raise subprocess.CalledProcessError(return_code, command)
|
||||||
|
return report
|
||||||
|
|
||||||
|
|
||||||
def uv_download_target():
|
def uv_download_target():
|
||||||
machine = platform.machine().lower()
|
machine = platform.machine().lower()
|
||||||
if sys.platform == "win32":
|
if sys.platform == "win32":
|
||||||
@@ -643,9 +711,15 @@ def install(args):
|
|||||||
installed_packages.extend(stage_installed)
|
installed_packages.extend(stage_installed)
|
||||||
runtime_offset = stage_end
|
runtime_offset = stage_end
|
||||||
write_status(args.status, VERIFY_SLICE_START, "verifying runtime")
|
write_status(args.status, VERIFY_SLICE_START, "verifying runtime")
|
||||||
run(verify_command(python))
|
device_report = run_verify(python, args.status)
|
||||||
write_status(args.status, VERIFY_SLICE_END, "Runtime is ready")
|
write_status(args.status, VERIFY_SLICE_END, "Runtime is ready")
|
||||||
|
|
||||||
|
devices = []
|
||||||
|
if device_report.get("cuda"):
|
||||||
|
devices.append("cuda")
|
||||||
|
if device_report.get("mps"):
|
||||||
|
devices.append("mps")
|
||||||
|
devices.append("cpu")
|
||||||
runtime_dir.mkdir(parents=True, exist_ok=True)
|
runtime_dir.mkdir(parents=True, exist_ok=True)
|
||||||
ready_path = runtime_dir / "snapform_ready.json"
|
ready_path = runtime_dir / "snapform_ready.json"
|
||||||
ready_path.write_text(
|
ready_path.write_text(
|
||||||
@@ -658,6 +732,7 @@ def install(args):
|
|||||||
"fastapi": FASTAPI_VERSION,
|
"fastapi": FASTAPI_VERSION,
|
||||||
"uvicorn": UVICORN_VERSION,
|
"uvicorn": UVICORN_VERSION,
|
||||||
"backend": args.backend,
|
"backend": args.backend,
|
||||||
|
"devices": devices,
|
||||||
"da3_commit": args.da3_commit,
|
"da3_commit": args.da3_commit,
|
||||||
"ben2_commit": args.ben2_commit,
|
"ben2_commit": args.ben2_commit,
|
||||||
"packages": sorted(set(installed_packages)),
|
"packages": sorted(set(installed_packages)),
|
||||||
|
|||||||
@@ -97,7 +97,7 @@ def extract_frame_range(source_object):
|
|||||||
)
|
)
|
||||||
frame_count = (
|
frame_count = (
|
||||||
min(image_frames, total_frames)
|
min(image_frames, total_frames)
|
||||||
if image_frames > 1
|
if image_frames >= 1
|
||||||
else total_frames
|
else total_frames
|
||||||
)
|
)
|
||||||
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
|
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
|
||||||
@@ -1400,12 +1400,14 @@ class RemoveImageBackground(bpy.types.Operator):
|
|||||||
@classmethod
|
@classmethod
|
||||||
def poll(cls, context):
|
def poll(cls, context):
|
||||||
from .operators import active_job
|
from .operators import active_job
|
||||||
from .runtime import host_busy
|
from .runtime import host_busy, model_ready
|
||||||
|
from .models import BEN2_MODEL
|
||||||
|
|
||||||
return (
|
return (
|
||||||
is_image_empty(context.object)
|
is_image_empty(context.object)
|
||||||
and active_job() is None
|
and active_job() is None
|
||||||
and not host_busy()
|
and not host_busy()
|
||||||
|
and model_ready(BEN2_MODEL.key)
|
||||||
)
|
)
|
||||||
|
|
||||||
def execute(self, context):
|
def execute(self, context):
|
||||||
@@ -1511,12 +1513,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
@classmethod
|
@classmethod
|
||||||
def poll(cls, context):
|
def poll(cls, context):
|
||||||
from .operators import active_job
|
from .operators import active_job
|
||||||
from .runtime import host_busy
|
from .runtime import host_busy, model_ready
|
||||||
|
from .models import MODELS
|
||||||
|
|
||||||
return (
|
return (
|
||||||
is_image_empty(context.object)
|
is_image_empty(context.object)
|
||||||
and active_job() is None
|
and active_job() is None
|
||||||
and not host_busy()
|
and not host_busy()
|
||||||
|
and any(model_ready(model_key) for model_key in MODELS)
|
||||||
)
|
)
|
||||||
|
|
||||||
def invoke(self, context, _event):
|
def invoke(self, context, _event):
|
||||||
@@ -1535,6 +1539,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
return context.window_manager.invoke_props_dialog(self, width=340)
|
return context.window_manager.invoke_props_dialog(self, width=340)
|
||||||
|
|
||||||
def draw(self, context):
|
def draw(self, context):
|
||||||
|
from .models import BEN2_MODEL
|
||||||
|
from .runtime import model_ready
|
||||||
|
|
||||||
layout = self.layout
|
layout = self.layout
|
||||||
layout.prop(self, "snapform_depth_source")
|
layout.prop(self, "snapform_depth_source")
|
||||||
if self.snapform_depth_source == "ESTIMATE":
|
if self.snapform_depth_source == "ESTIMATE":
|
||||||
@@ -1563,7 +1570,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
|
|||||||
if settings_body is not None:
|
if settings_body is not None:
|
||||||
settings_body.prop(self, "snapform_displacement_via")
|
settings_body.prop(self, "snapform_displacement_via")
|
||||||
settings_body.prop(self, "snapform_subdivision_level")
|
settings_body.prop(self, "snapform_subdivision_level")
|
||||||
settings_body.prop(self, "snapform_remove_background")
|
remove_background_row = settings_body.row()
|
||||||
|
remove_background_row.enabled = model_ready(BEN2_MODEL.key)
|
||||||
|
remove_background_row.prop(self, "snapform_remove_background")
|
||||||
settings_body.prop(self, "snapform_unshaded")
|
settings_body.prop(self, "snapform_unshaded")
|
||||||
layout.prop(self, "snapform_delete_original")
|
layout.prop(self, "snapform_delete_original")
|
||||||
|
|
||||||
|
|||||||
@@ -5,6 +5,7 @@ from dataclasses import dataclass
|
|||||||
class ModelSpec:
|
class ModelSpec:
|
||||||
key: str
|
key: str
|
||||||
label: str
|
label: str
|
||||||
|
directory_name: str
|
||||||
repository: str
|
repository: str
|
||||||
license_name: str
|
license_name: str
|
||||||
description: str
|
description: str
|
||||||
@@ -16,6 +17,7 @@ MODELS = {
|
|||||||
"SMALL": ModelSpec(
|
"SMALL": ModelSpec(
|
||||||
key="SMALL",
|
key="SMALL",
|
||||||
label="DA3 Small",
|
label="DA3 Small",
|
||||||
|
directory_name="da3_small",
|
||||||
repository="depth-anything/DA3-SMALL",
|
repository="depth-anything/DA3-SMALL",
|
||||||
license_name="Apache 2.0",
|
license_name="Apache 2.0",
|
||||||
description="Fastest any-view model with Apache 2.0 weights",
|
description="Fastest any-view model with Apache 2.0 weights",
|
||||||
@@ -24,6 +26,7 @@ MODELS = {
|
|||||||
"BASE": ModelSpec(
|
"BASE": ModelSpec(
|
||||||
key="BASE",
|
key="BASE",
|
||||||
label="DA3 Base",
|
label="DA3 Base",
|
||||||
|
directory_name="da3_base",
|
||||||
repository="depth-anything/DA3-BASE",
|
repository="depth-anything/DA3-BASE",
|
||||||
license_name="Apache 2.0",
|
license_name="Apache 2.0",
|
||||||
description="Balanced any-view model with Apache 2.0 weights",
|
description="Balanced any-view model with Apache 2.0 weights",
|
||||||
@@ -32,6 +35,7 @@ MODELS = {
|
|||||||
"LARGE_1_1": ModelSpec(
|
"LARGE_1_1": ModelSpec(
|
||||||
key="LARGE_1_1",
|
key="LARGE_1_1",
|
||||||
label="DA3 Large 1.1",
|
label="DA3 Large 1.1",
|
||||||
|
directory_name="da3_large_1_1",
|
||||||
repository="depth-anything/DA3-LARGE-1.1",
|
repository="depth-anything/DA3-LARGE-1.1",
|
||||||
license_name="CC BY-NC 4.0",
|
license_name="CC BY-NC 4.0",
|
||||||
description="Higher quality non-commercial model with larger VRAM use",
|
description="Higher quality non-commercial model with larger VRAM use",
|
||||||
@@ -40,6 +44,7 @@ MODELS = {
|
|||||||
"GIANT_1_1": ModelSpec(
|
"GIANT_1_1": ModelSpec(
|
||||||
key="GIANT_1_1",
|
key="GIANT_1_1",
|
||||||
label="DA3 Giant 1.1",
|
label="DA3 Giant 1.1",
|
||||||
|
directory_name="da3_giant_1_1",
|
||||||
repository="depth-anything/DA3-GIANT-1.1",
|
repository="depth-anything/DA3-GIANT-1.1",
|
||||||
license_name="CC BY-NC 4.0",
|
license_name="CC BY-NC 4.0",
|
||||||
description="Non-commercial Giant model with Gaussian head",
|
description="Non-commercial Giant model with Gaussian head",
|
||||||
@@ -48,6 +53,7 @@ MODELS = {
|
|||||||
"NESTED_GIANT_LARGE_1_1": ModelSpec(
|
"NESTED_GIANT_LARGE_1_1": ModelSpec(
|
||||||
key="NESTED_GIANT_LARGE_1_1",
|
key="NESTED_GIANT_LARGE_1_1",
|
||||||
label="DA3 Nested Giant Large 1.1",
|
label="DA3 Nested Giant Large 1.1",
|
||||||
|
directory_name="da3_nested_giant_large_1_1",
|
||||||
repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
|
repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
|
||||||
license_name="CC BY-NC 4.0",
|
license_name="CC BY-NC 4.0",
|
||||||
description="Non-commercial nested model with metric depth and Gaussian head",
|
description="Non-commercial nested model with metric depth and Gaussian head",
|
||||||
@@ -58,6 +64,7 @@ MODELS = {
|
|||||||
BEN2_MODEL = ModelSpec(
|
BEN2_MODEL = ModelSpec(
|
||||||
key="BEN2_BASE",
|
key="BEN2_BASE",
|
||||||
label="BEN2 Base",
|
label="BEN2 Base",
|
||||||
|
directory_name="ben2_base",
|
||||||
repository="PramaLLC/BEN2",
|
repository="PramaLLC/BEN2",
|
||||||
license_name="MIT",
|
license_name="MIT",
|
||||||
description="Fast high-quality background removal model",
|
description="Fast high-quality background removal model",
|
||||||
@@ -142,3 +149,8 @@ def get_downloadable_model(model_key):
|
|||||||
return DOWNLOADABLE_MODELS[model_key]
|
return DOWNLOADABLE_MODELS[model_key]
|
||||||
except KeyError as error:
|
except KeyError as error:
|
||||||
raise ValueError(f"Unknown model: {model_key}") from error
|
raise ValueError(f"Unknown model: {model_key}") from error
|
||||||
|
|
||||||
|
|
||||||
|
def model_directory_name(model_key):
|
||||||
|
"""Return the model's folder name inside the shared models directory."""
|
||||||
|
return get_downloadable_model(model_key).directory_name
|
||||||
|
|||||||
@@ -337,6 +337,7 @@ class RunJob(bpy.types.Operator):
|
|||||||
paths["result"],
|
paths["result"],
|
||||||
BEN2_MODEL.key,
|
BEN2_MODEL.key,
|
||||||
start_frame=start_frame,
|
start_frame=start_frame,
|
||||||
|
device=self._task_device(settings),
|
||||||
),
|
),
|
||||||
paths,
|
paths,
|
||||||
)
|
)
|
||||||
@@ -398,10 +399,21 @@ class RunJob(bpy.types.Operator):
|
|||||||
self.job_type == IMAGE_DEPTH_JOB
|
self.job_type == IMAGE_DEPTH_JOB
|
||||||
and remove_background
|
and remove_background
|
||||||
),
|
),
|
||||||
|
device=self._task_device(settings),
|
||||||
),
|
),
|
||||||
paths,
|
paths,
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _task_device(self, settings):
|
||||||
|
"""Debug panels use the scene device; production tasks use Preferences."""
|
||||||
|
if self.job_type in {
|
||||||
|
IMAGE_DEPTH_JOB,
|
||||||
|
IMAGE_BACKGROUND_JOB,
|
||||||
|
ORIGINAL_BACKGROUND_JOB,
|
||||||
|
}:
|
||||||
|
return runtime.production_device()
|
||||||
|
return settings.device
|
||||||
|
|
||||||
def _require_online_access(self):
|
def _require_online_access(self):
|
||||||
if not bpy.app.online_access:
|
if not bpy.app.online_access:
|
||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
|
|||||||
@@ -1,9 +1,9 @@
|
|||||||
import shutil
|
|
||||||
import sys
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
|
|
||||||
import bpy
|
import bpy
|
||||||
|
|
||||||
|
from .runtime import device_enum_items
|
||||||
|
|
||||||
|
|
||||||
class SnapFormPreferences(bpy.types.AddonPreferences):
|
class SnapFormPreferences(bpy.types.AddonPreferences):
|
||||||
bl_idname = __package__
|
bl_idname = __package__
|
||||||
@@ -17,31 +17,14 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
|
|||||||
subtype="DIR_PATH",
|
subtype="DIR_PATH",
|
||||||
default=str(Path.home() / ".snapform"),
|
default=str(Path.home() / ".snapform"),
|
||||||
)
|
)
|
||||||
uv_executable: bpy.props.StringProperty(
|
device: bpy.props.EnumProperty(
|
||||||
name="uv Executable",
|
name="Device",
|
||||||
description=(
|
description=(
|
||||||
"Path to uv; leave empty to auto-download a standalone uv "
|
"Inference device used by production tasks; only devices supported "
|
||||||
"when it is missing"
|
"on this machine are listed"
|
||||||
),
|
),
|
||||||
subtype="FILE_PATH",
|
items=device_enum_items,
|
||||||
default=shutil.which("uv") or "",
|
default=0,
|
||||||
)
|
|
||||||
torch_backend: bpy.props.EnumProperty(
|
|
||||||
name="Torch Backend",
|
|
||||||
description="PyTorch build installed into the isolated Runtime",
|
|
||||||
items=(
|
|
||||||
("CU128", "CUDA 12.8", "NVIDIA CUDA 12.8 PyTorch wheels"),
|
|
||||||
("CPU", "CPU", "CPU-only PyTorch wheels"),
|
|
||||||
),
|
|
||||||
default="CPU" if sys.platform == "darwin" else "CU128",
|
|
||||||
)
|
|
||||||
pypi_index_url: bpy.props.StringProperty(
|
|
||||||
name="PyPI Index",
|
|
||||||
description=(
|
|
||||||
"Package index for the isolated Runtime; leave empty to auto-detect "
|
|
||||||
"with mirror fallback. Example: https://pypi.tuna.tsinghua.edu.cn/simple"
|
|
||||||
),
|
|
||||||
default="",
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def draw(self, _context):
|
def draw(self, _context):
|
||||||
@@ -50,9 +33,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
|
|||||||
|
|
||||||
layout = self.layout
|
layout = self.layout
|
||||||
layout.prop(self, "storage_root")
|
layout.prop(self, "storage_root")
|
||||||
layout.prop(self, "uv_executable")
|
layout.prop(self, "device")
|
||||||
layout.prop(self, "torch_backend")
|
|
||||||
layout.prop(self, "pypi_index_url")
|
|
||||||
layout.separator()
|
layout.separator()
|
||||||
layout.label(text=f"Runtime: {runtime_directory()}")
|
layout.label(text=f"Runtime: {runtime_directory()}")
|
||||||
if runtime_ready():
|
if runtime_ready():
|
||||||
|
|||||||
@@ -1,6 +1,7 @@
|
|||||||
import bpy
|
import bpy
|
||||||
|
|
||||||
from .models import gs_model_enum_items, model_enum_items
|
from .models import gs_model_enum_items, model_enum_items
|
||||||
|
from .runtime import device_enum_items
|
||||||
|
|
||||||
|
|
||||||
class SnapFormSettings(bpy.types.PropertyGroup):
|
class SnapFormSettings(bpy.types.PropertyGroup):
|
||||||
@@ -18,7 +19,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
|
|||||||
name="Extract Frames",
|
name="Extract Frames",
|
||||||
description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
|
description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
|
||||||
default=16,
|
default=16,
|
||||||
min=2,
|
min=1,
|
||||||
soft_max=64,
|
soft_max=64,
|
||||||
)
|
)
|
||||||
model: bpy.props.EnumProperty(
|
model: bpy.props.EnumProperty(
|
||||||
@@ -32,15 +33,13 @@ class SnapFormSettings(bpy.types.PropertyGroup):
|
|||||||
items=gs_model_enum_items,
|
items=gs_model_enum_items,
|
||||||
)
|
)
|
||||||
device: bpy.props.EnumProperty(
|
device: bpy.props.EnumProperty(
|
||||||
name="Device",
|
name="Debug Device",
|
||||||
description="Inference device used by the external Runtime",
|
description=(
|
||||||
items=(
|
"Inference device used by Debug tasks only; production tasks use "
|
||||||
("AUTO", "Auto", "Use CUDA, then MPS, otherwise CPU"),
|
"the Device setting in the Add-on Preferences"
|
||||||
("CUDA", "CUDA", "Require an NVIDIA CUDA device"),
|
|
||||||
("MPS", "MPS", "Use Apple Metal (MPS)"),
|
|
||||||
("CPU", "CPU", "Run on CPU"),
|
|
||||||
),
|
),
|
||||||
default="AUTO",
|
items=device_enum_items,
|
||||||
|
default=0,
|
||||||
)
|
)
|
||||||
process_res: bpy.props.IntProperty(
|
process_res: bpy.props.IntProperty(
|
||||||
name="Process Res",
|
name="Process Res",
|
||||||
|
|||||||
@@ -1,5 +1,6 @@
|
|||||||
import json
|
import json
|
||||||
import os
|
import os
|
||||||
|
import platform
|
||||||
import secrets
|
import secrets
|
||||||
import shutil
|
import shutil
|
||||||
import socket
|
import socket
|
||||||
@@ -15,7 +16,12 @@ from urllib.request import ProxyHandler, Request, build_opener
|
|||||||
|
|
||||||
import bpy
|
import bpy
|
||||||
|
|
||||||
from .models import BEN2_MODEL, get_downloadable_model, selected_model_key
|
from .models import (
|
||||||
|
BEN2_MODEL,
|
||||||
|
get_downloadable_model,
|
||||||
|
model_directory_name,
|
||||||
|
selected_model_key,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
|
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
|
||||||
@@ -107,7 +113,16 @@ def models_directory():
|
|||||||
|
|
||||||
|
|
||||||
def model_directory(model_key):
|
def model_directory(model_key):
|
||||||
return models_directory() / model_key.lower()
|
root = models_directory()
|
||||||
|
directory = root / model_directory_name(model_key)
|
||||||
|
legacy_directory = root / model_key.lower()
|
||||||
|
if (
|
||||||
|
not directory.is_dir()
|
||||||
|
and legacy_directory.is_dir()
|
||||||
|
and legacy_directory != directory
|
||||||
|
):
|
||||||
|
legacy_directory.rename(directory)
|
||||||
|
return directory
|
||||||
|
|
||||||
|
|
||||||
def model_ready(model_key):
|
def model_ready(model_key):
|
||||||
@@ -117,6 +132,61 @@ def model_ready(model_key):
|
|||||||
)
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def runtime_devices():
|
||||||
|
"""Devices detected by the installed Runtime, or None when unknown."""
|
||||||
|
ready_path = runtime_directory() / "snapform_ready.json"
|
||||||
|
if not ready_path.is_file():
|
||||||
|
return None
|
||||||
|
try:
|
||||||
|
payload = json.loads(ready_path.read_text(encoding="utf-8"))
|
||||||
|
except (json.JSONDecodeError, OSError):
|
||||||
|
return None
|
||||||
|
devices = payload.get("devices")
|
||||||
|
if not isinstance(devices, list) or not devices:
|
||||||
|
return None
|
||||||
|
return [str(name).lower() for name in devices]
|
||||||
|
|
||||||
|
|
||||||
|
def machine_devices():
|
||||||
|
"""Devices the machine can run, fastest first; used before Runtime install."""
|
||||||
|
devices = ["cpu"]
|
||||||
|
if shutil.which("nvidia-smi"):
|
||||||
|
devices.insert(0, "cuda")
|
||||||
|
elif sys.platform == "darwin" and platform.machine().lower() in {
|
||||||
|
"arm64",
|
||||||
|
"aarch64",
|
||||||
|
}:
|
||||||
|
devices.insert(0, "mps")
|
||||||
|
return devices
|
||||||
|
|
||||||
|
|
||||||
|
def supported_devices():
|
||||||
|
"""Inference devices available on this machine, fastest first."""
|
||||||
|
return runtime_devices() or machine_devices()
|
||||||
|
|
||||||
|
|
||||||
|
DEVICE_LABELS = {
|
||||||
|
"cuda": ("CUDA", "NVIDIA CUDA"),
|
||||||
|
"mps": ("MPS", "Apple Metal (MPS)"),
|
||||||
|
"cpu": ("CPU", "CPU"),
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def device_enum_items(_owner, _context):
|
||||||
|
"""Blender EnumProperty items listing only devices supported on this machine."""
|
||||||
|
return [
|
||||||
|
(name.upper(), *DEVICE_LABELS[name], index)
|
||||||
|
for index, name in enumerate(supported_devices())
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def production_device():
|
||||||
|
"""Device configured for production tasks, falling back to Auto."""
|
||||||
|
addon_preferences = preferences()
|
||||||
|
configured = getattr(addon_preferences, "device", "")
|
||||||
|
return configured or "AUTO"
|
||||||
|
|
||||||
|
|
||||||
def jobs_directory():
|
def jobs_directory():
|
||||||
path = storage_root() / "jobs"
|
path = storage_root() / "jobs"
|
||||||
path.mkdir(parents=True, exist_ok=True)
|
path.mkdir(parents=True, exist_ok=True)
|
||||||
@@ -153,13 +223,7 @@ def new_job_paths(job_type):
|
|||||||
|
|
||||||
|
|
||||||
def uv_executable():
|
def uv_executable():
|
||||||
addon_preferences = preferences()
|
"""Return the bundled uv path inside the shared storage."""
|
||||||
configured = getattr(addon_preferences, "uv_executable", "uv")
|
|
||||||
resolved = shutil.which(configured)
|
|
||||||
if resolved:
|
|
||||||
return resolved
|
|
||||||
if Path(configured).is_file():
|
|
||||||
return str(Path(configured))
|
|
||||||
return str(uv_cache_path())
|
return str(uv_cache_path())
|
||||||
|
|
||||||
|
|
||||||
@@ -171,9 +235,7 @@ def uv_cache_path():
|
|||||||
|
|
||||||
def install_runtime_command(status_path):
|
def install_runtime_command(status_path):
|
||||||
executable = uv_executable()
|
executable = uv_executable()
|
||||||
addon_preferences = preferences()
|
backend = "cu128" if "cuda" in machine_devices() else "cpu"
|
||||||
backend = getattr(addon_preferences, "torch_backend", "CU128")
|
|
||||||
pypi_index = getattr(addon_preferences, "pypi_index_url", "") or ""
|
|
||||||
script = Path(__file__).with_name("bootstrap_runtime.py")
|
script = Path(__file__).with_name("bootstrap_runtime.py")
|
||||||
command = [
|
command = [
|
||||||
sys.executable,
|
sys.executable,
|
||||||
@@ -185,7 +247,7 @@ def install_runtime_command(status_path):
|
|||||||
"--models-dir",
|
"--models-dir",
|
||||||
str(models_directory()),
|
str(models_directory()),
|
||||||
"--backend",
|
"--backend",
|
||||||
backend.lower(),
|
backend,
|
||||||
"--da3-commit",
|
"--da3-commit",
|
||||||
DA3_COMMIT,
|
DA3_COMMIT,
|
||||||
"--ben2-commit",
|
"--ben2-commit",
|
||||||
@@ -193,8 +255,6 @@ def install_runtime_command(status_path):
|
|||||||
"--status",
|
"--status",
|
||||||
str(status_path),
|
str(status_path),
|
||||||
]
|
]
|
||||||
if pypi_index:
|
|
||||||
command.extend(["--pypi-index", pypi_index])
|
|
||||||
return command
|
return command
|
||||||
|
|
||||||
|
|
||||||
@@ -640,6 +700,7 @@ def remove_background_request(
|
|||||||
model_key,
|
model_key,
|
||||||
video_frames=None,
|
video_frames=None,
|
||||||
start_frame=1,
|
start_frame=1,
|
||||||
|
device=None,
|
||||||
):
|
):
|
||||||
return HostJobRequest(
|
return HostJobRequest(
|
||||||
{
|
{
|
||||||
@@ -650,7 +711,7 @@ def remove_background_request(
|
|||||||
"model_dir": str(model_directory(model_key)),
|
"model_dir": str(model_directory(model_key)),
|
||||||
"output_dir": str(result_directory),
|
"output_dir": str(result_directory),
|
||||||
"status": str(status_path),
|
"status": str(status_path),
|
||||||
"device": settings.device.lower(),
|
"device": (device or settings.device).lower(),
|
||||||
"video_frames": (
|
"video_frames": (
|
||||||
settings.video_frames
|
settings.video_frames
|
||||||
if video_frames is None
|
if video_frames is None
|
||||||
@@ -701,6 +762,7 @@ def reconstruct_request(
|
|||||||
use_ray_pose=None,
|
use_ray_pose=None,
|
||||||
remove_background=False,
|
remove_background=False,
|
||||||
start_frame=1,
|
start_frame=1,
|
||||||
|
device=None,
|
||||||
):
|
):
|
||||||
snapform_selected_key = model_key or selected_model_key(settings, export_format)
|
snapform_selected_key = model_key or selected_model_key(settings, export_format)
|
||||||
snapform_selected_process_res = (
|
snapform_selected_process_res = (
|
||||||
@@ -732,7 +794,7 @@ def reconstruct_request(
|
|||||||
"gs_views_interval": settings.gs_views_interval,
|
"gs_views_interval": settings.gs_views_interval,
|
||||||
"video_frames": settings.video_frames,
|
"video_frames": settings.video_frames,
|
||||||
"start_frame": start_frame,
|
"start_frame": start_frame,
|
||||||
"device": settings.device.lower(),
|
"device": (device or settings.device).lower(),
|
||||||
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
|
"use_ray_pose": 1 if snapform_selected_ray_pose else 0,
|
||||||
"ref_view_strategy": settings.ref_view_strategy,
|
"ref_view_strategy": settings.ref_view_strategy,
|
||||||
"remove_background_model_dir": "",
|
"remove_background_model_dir": "",
|
||||||
|
|||||||
@@ -1356,6 +1356,104 @@ class AddonTest(unittest.TestCase):
|
|||||||
["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"],
|
["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"],
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_device_enums_list_only_supported_devices(self):
|
||||||
|
runtime = importlib.import_module("snapform.runtime")
|
||||||
|
with patch.object(
|
||||||
|
runtime,
|
||||||
|
"supported_devices",
|
||||||
|
return_value=["cuda", "cpu"],
|
||||||
|
):
|
||||||
|
items = runtime.device_enum_items(None, None)
|
||||||
|
|
||||||
|
self.assertEqual([item[0] for item in items], ["CUDA", "CPU"])
|
||||||
|
self.assertEqual(items[0][2], "NVIDIA CUDA")
|
||||||
|
|
||||||
|
def test_device_enums_default_to_the_fastest_supported_device(self):
|
||||||
|
for class_type in (
|
||||||
|
self.addon.SnapFormSettings,
|
||||||
|
self.addon.SnapFormPreferences,
|
||||||
|
):
|
||||||
|
with self.subTest(class_type=class_type.__name__):
|
||||||
|
self.assertEqual(
|
||||||
|
class_type.__annotations__["device"]["default"],
|
||||||
|
0,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_preferences_expose_production_device(self):
|
||||||
|
preferences = self.addon.SnapFormPreferences
|
||||||
|
|
||||||
|
self.assertIn("device", preferences.__annotations__)
|
||||||
|
items_callback = preferences.__annotations__["device"]["items"]
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"supported_devices",
|
||||||
|
return_value=["cpu"],
|
||||||
|
):
|
||||||
|
items = items_callback(None, None)
|
||||||
|
self.assertEqual([item[0] for item in items], ["CPU"])
|
||||||
|
|
||||||
|
def test_remove_background_poll_requires_ben2_model(self):
|
||||||
|
context = SimpleNamespace(
|
||||||
|
object=SimpleNamespace(
|
||||||
|
type="EMPTY",
|
||||||
|
empty_display_type="IMAGE",
|
||||||
|
data=object(),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"host_busy",
|
||||||
|
return_value=False,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
return_value=False,
|
||||||
|
):
|
||||||
|
self.assertFalse(self.addon.RemoveImageBackground.poll(context))
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"host_busy",
|
||||||
|
return_value=False,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
return_value=True,
|
||||||
|
):
|
||||||
|
self.assertTrue(self.addon.RemoveImageBackground.poll(context))
|
||||||
|
|
||||||
|
def test_convert_to_depth_plane_poll_requires_da3_model(self):
|
||||||
|
context = SimpleNamespace(
|
||||||
|
object=SimpleNamespace(
|
||||||
|
type="EMPTY",
|
||||||
|
empty_display_type="IMAGE",
|
||||||
|
data=object(),
|
||||||
|
)
|
||||||
|
)
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"host_busy",
|
||||||
|
return_value=False,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
side_effect=lambda key: key == "BASE",
|
||||||
|
):
|
||||||
|
self.assertTrue(
|
||||||
|
self.addon.ConvertImageToDepthPlane.poll(context)
|
||||||
|
)
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"host_busy",
|
||||||
|
return_value=False,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
return_value=False,
|
||||||
|
):
|
||||||
|
self.assertFalse(
|
||||||
|
self.addon.ConvertImageToDepthPlane.poll(context)
|
||||||
|
)
|
||||||
|
|
||||||
def test_debug_panel_is_closed_by_default(self):
|
def test_debug_panel_is_closed_by_default(self):
|
||||||
self.assertIn(
|
self.assertIn(
|
||||||
"DEFAULT_CLOSED",
|
"DEFAULT_CLOSED",
|
||||||
@@ -1485,6 +1583,95 @@ class AddonTest(unittest.TestCase):
|
|||||||
self.assertEqual(remove_background.call_args.args[1], input_path)
|
self.assertEqual(remove_background.call_args.args[1], input_path)
|
||||||
self.assertEqual(remove_background.call_args.args[4], "BEN2_BASE")
|
self.assertEqual(remove_background.call_args.args[4], "BEN2_BASE")
|
||||||
|
|
||||||
|
def test_debug_jobs_use_the_scene_device(self):
|
||||||
|
settings = SimpleNamespace(
|
||||||
|
background_input_path="input.png",
|
||||||
|
device="CPU",
|
||||||
|
)
|
||||||
|
operator = self.addon.RunJob()
|
||||||
|
operator.job_type = self.addon.operators.BACKGROUND_REMOVAL_JOB
|
||||||
|
|
||||||
|
with TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
input_path = root / "input.png"
|
||||||
|
input_path.touch()
|
||||||
|
settings.background_input_path = str(input_path)
|
||||||
|
paths = {
|
||||||
|
"log": root / "job.log",
|
||||||
|
"status": root / "job.json",
|
||||||
|
"result": root / "result",
|
||||||
|
}
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"new_job_paths",
|
||||||
|
return_value=paths,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"runtime_ready",
|
||||||
|
return_value=True,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
return_value=True,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"remove_background_request",
|
||||||
|
return_value=self.addon.runtime.HostJobRequest(
|
||||||
|
{"job_id": "test"}
|
||||||
|
),
|
||||||
|
) as remove_background:
|
||||||
|
operator._prepare_command(None, settings)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
remove_background.call_args.kwargs["device"],
|
||||||
|
"CPU",
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_production_jobs_use_the_preferences_device(self):
|
||||||
|
settings = SimpleNamespace(device="CPU", video_frames=16)
|
||||||
|
operator = self.addon.RunJob()
|
||||||
|
operator.job_type = self.addon.operators.IMAGE_BACKGROUND_JOB
|
||||||
|
|
||||||
|
with TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
input_path = root / "input.png"
|
||||||
|
input_path.touch()
|
||||||
|
operator.snapform_input_path = str(input_path)
|
||||||
|
paths = {
|
||||||
|
"log": root / "job.log",
|
||||||
|
"status": root / "job.json",
|
||||||
|
"result": root / "result",
|
||||||
|
}
|
||||||
|
with patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"new_job_paths",
|
||||||
|
return_value=paths,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"runtime_ready",
|
||||||
|
return_value=True,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"model_ready",
|
||||||
|
return_value=True,
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"production_device",
|
||||||
|
return_value="MPS",
|
||||||
|
), patch.object(
|
||||||
|
self.addon.runtime,
|
||||||
|
"remove_background_request",
|
||||||
|
return_value=self.addon.runtime.HostJobRequest(
|
||||||
|
{"job_id": "test"}
|
||||||
|
),
|
||||||
|
) as remove_background:
|
||||||
|
operator._prepare_command(None, settings)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
remove_background.call_args.kwargs["device"],
|
||||||
|
"MPS",
|
||||||
|
)
|
||||||
|
|
||||||
def test_image_background_job_uses_operator_input(self):
|
def test_image_background_job_uses_operator_input(self):
|
||||||
settings = SimpleNamespace(device="AUTO")
|
settings = SimpleNamespace(device="AUTO")
|
||||||
operator = self.addon.RunJob()
|
operator = self.addon.RunJob()
|
||||||
@@ -1804,8 +1991,8 @@ class AddonTest(unittest.TestCase):
|
|||||||
return_value="C:/uv/uv.exe",
|
return_value="C:/uv/uv.exe",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
runtime,
|
runtime,
|
||||||
"preferences",
|
"machine_devices",
|
||||||
return_value=SimpleNamespace(torch_backend="CU128"),
|
return_value=["cuda", "cpu"],
|
||||||
), patch.object(
|
), patch.object(
|
||||||
runtime,
|
runtime,
|
||||||
"runtime_directory",
|
"runtime_directory",
|
||||||
@@ -1821,11 +2008,15 @@ class AddonTest(unittest.TestCase):
|
|||||||
self.assertEqual(command[0], sys.executable)
|
self.assertEqual(command[0], sys.executable)
|
||||||
self.assertIn("--uv", command)
|
self.assertIn("--uv", command)
|
||||||
self.assertIn("C:/uv/uv.exe", command)
|
self.assertIn("C:/uv/uv.exe", command)
|
||||||
|
self.assertEqual(
|
||||||
|
command[command.index("--backend") + 1],
|
||||||
|
"cu128",
|
||||||
|
)
|
||||||
self.assertTrue(
|
self.assertTrue(
|
||||||
any("bootstrap_runtime.py" in str(value) for value in command)
|
any("bootstrap_runtime.py" in str(value) for value in command)
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_install_runtime_command_passes_configured_pypi_index(self):
|
def test_install_runtime_command_auto_selects_backend(self):
|
||||||
runtime = importlib.import_module("snapform.runtime")
|
runtime = importlib.import_module("snapform.runtime")
|
||||||
with patch.object(
|
with patch.object(
|
||||||
runtime,
|
runtime,
|
||||||
@@ -1833,11 +2024,8 @@ class AddonTest(unittest.TestCase):
|
|||||||
return_value="C:/uv/uv.exe",
|
return_value="C:/uv/uv.exe",
|
||||||
), patch.object(
|
), patch.object(
|
||||||
runtime,
|
runtime,
|
||||||
"preferences",
|
"machine_devices",
|
||||||
return_value=SimpleNamespace(
|
return_value=["cpu"],
|
||||||
torch_backend="CU128",
|
|
||||||
pypi_index_url="https://pypi.tuna.tsinghua.edu.cn/simple",
|
|
||||||
),
|
|
||||||
), patch.object(
|
), patch.object(
|
||||||
runtime,
|
runtime,
|
||||||
"runtime_directory",
|
"runtime_directory",
|
||||||
@@ -1849,11 +2037,12 @@ class AddonTest(unittest.TestCase):
|
|||||||
):
|
):
|
||||||
command = runtime.install_runtime_command(Path("status.json"))
|
command = runtime.install_runtime_command(Path("status.json"))
|
||||||
|
|
||||||
index_position = command.index("--pypi-index")
|
self.assertNotIn("--pypi-index", command)
|
||||||
self.assertEqual(
|
self.assertEqual(
|
||||||
command[index_position + 1],
|
command[command.index("--backend") + 1],
|
||||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
"cpu",
|
||||||
)
|
)
|
||||||
|
self.assertIn("C:/uv/uv.exe", command)
|
||||||
|
|
||||||
def test_completed_export_opens_job_directory(self):
|
def test_completed_export_opens_job_directory(self):
|
||||||
calls = []
|
calls = []
|
||||||
|
|||||||
@@ -109,7 +109,7 @@ class HostTest(unittest.TestCase):
|
|||||||
|
|
||||||
self.assertIs(first, second)
|
self.assertIs(first, second)
|
||||||
self.assertEqual(loads, ["model"])
|
self.assertEqual(loads, ["model"])
|
||||||
self.assertEqual(cache.status()["da3"], "da3_model")
|
self.assertEqual(cache.status()["da3"], "model")
|
||||||
|
|
||||||
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
|
def test_host_tracks_clients_and_rejects_shared_shutdown(self):
|
||||||
fake_models = SimpleNamespace(
|
fake_models = SimpleNamespace(
|
||||||
|
|||||||
@@ -62,7 +62,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
) as mocked_run, patch.object(
|
) as mocked_run, patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
) as mocked_run_install:
|
) as mocked_run_install, patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
commands = [call.args[0] for call in mocked_run.call_args_list]
|
commands = [call.args[0] for call in mocked_run.call_args_list]
|
||||||
@@ -113,10 +117,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
|
self.assertIn(f"uvicorn=={bootstrap.UVICORN_VERSION}", host_command)
|
||||||
self.assertIn("setuptools", host_command)
|
self.assertIn("setuptools", host_command)
|
||||||
self.assertNotIn("ninja", host_command)
|
self.assertNotIn("ninja", host_command)
|
||||||
check_command = next(
|
check_command = bootstrap.verify_command(
|
||||||
command
|
bootstrap.environment_python(runtime_dir)
|
||||||
for command in commands
|
)
|
||||||
if any("import torch;" in str(value) for value in command)
|
self.assertTrue(
|
||||||
|
any("import torch;" in str(value) for value in check_command)
|
||||||
)
|
)
|
||||||
self.assertTrue(
|
self.assertTrue(
|
||||||
any(
|
any(
|
||||||
@@ -130,6 +135,12 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
self.assertFalse(
|
self.assertFalse(
|
||||||
any("import gsplat" in str(value) for value in check_command)
|
any("import gsplat" in str(value) for value in check_command)
|
||||||
)
|
)
|
||||||
|
self.assertTrue(
|
||||||
|
any(
|
||||||
|
"SNAPFORM_DEVICES=" in str(value)
|
||||||
|
for value in check_command
|
||||||
|
)
|
||||||
|
)
|
||||||
default_downloads = [
|
default_downloads = [
|
||||||
command
|
command
|
||||||
for command in commands
|
for command in commands
|
||||||
@@ -190,7 +201,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
) as mocked_run_install:
|
) as mocked_run_install, patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
commands = [
|
commands = [
|
||||||
@@ -372,7 +387,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
bootstrap,
|
bootstrap,
|
||||||
"run_stage",
|
"run_stage",
|
||||||
side_effect=fake_run_stage,
|
side_effect=fake_run_stage,
|
||||||
), patch.object(bootstrap, "run"):
|
), patch.object(bootstrap, "run"), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
starts = [item[1] for item in stage_ranges]
|
starts = [item[1] for item in stage_ranges]
|
||||||
@@ -604,34 +623,6 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
"https://pypi.tuna.tsinghua.edu.cn/simple",
|
||||||
)
|
)
|
||||||
|
|
||||||
def test_child_environment_injects_registry_proxy(self):
|
|
||||||
bootstrap = load_bootstrap_runtime()
|
|
||||||
with patch.object(
|
|
||||||
bootstrap.urllib.request,
|
|
||||||
"getproxies",
|
|
||||||
return_value={
|
|
||||||
"http": "http://proxy:8080",
|
|
||||||
"https": "http://proxy:8080",
|
|
||||||
},
|
|
||||||
):
|
|
||||||
env = bootstrap.child_environment()
|
|
||||||
|
|
||||||
self.assertEqual(env["HTTP_PROXY"], "http://proxy:8080")
|
|
||||||
self.assertEqual(env["HTTPS_PROXY"], "http://proxy:8080")
|
|
||||||
|
|
||||||
def test_child_environment_keeps_explicit_proxy(self):
|
|
||||||
bootstrap = load_bootstrap_runtime()
|
|
||||||
with patch.object(
|
|
||||||
bootstrap.urllib.request,
|
|
||||||
"getproxies",
|
|
||||||
return_value={"https": "http://proxy:8080"},
|
|
||||||
):
|
|
||||||
env = bootstrap.child_environment(
|
|
||||||
{"HTTPS_PROXY": "http://explicit:3128"}
|
|
||||||
)
|
|
||||||
|
|
||||||
self.assertEqual(env["HTTPS_PROXY"], "http://explicit:3128")
|
|
||||||
|
|
||||||
def test_ready_json_records_installed_packages(self):
|
def test_ready_json_records_installed_packages(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
@@ -657,6 +648,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
bootstrap,
|
bootstrap,
|
||||||
"run_stage",
|
"run_stage",
|
||||||
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
|
return_value=["torch==2.9.1+cu128", "numpy==1.26.4"],
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={"cuda": True, "mps": False},
|
||||||
):
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
@@ -669,6 +664,7 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
ready["packages"],
|
ready["packages"],
|
||||||
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
["numpy==1.26.4", "torch==2.9.1+cu128"],
|
||||||
)
|
)
|
||||||
|
self.assertEqual(ready["devices"], ["cuda", "cpu"])
|
||||||
|
|
||||||
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
def test_broken_runtime_is_cleaned_before_reinstall(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
@@ -693,6 +689,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
):
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
@@ -728,6 +728,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
):
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
@@ -758,7 +762,11 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
bootstrap,
|
bootstrap,
|
||||||
"run",
|
"run",
|
||||||
side_effect=fake_run,
|
side_effect=fake_run,
|
||||||
), patch.object(bootstrap, "run_install"):
|
), patch.object(bootstrap, "run_install"), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
self.assertTrue((runtime_dir / "snapform_ready.json").is_file())
|
||||||
@@ -872,6 +880,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
side_effect=fake_run_install,
|
side_effect=fake_run_install,
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
):
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
@@ -918,6 +930,10 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
), patch.object(
|
), patch.object(
|
||||||
bootstrap,
|
bootstrap,
|
||||||
"run_install",
|
"run_install",
|
||||||
|
), patch.object(
|
||||||
|
bootstrap,
|
||||||
|
"run_verify",
|
||||||
|
return_value={},
|
||||||
):
|
):
|
||||||
bootstrap.install(args)
|
bootstrap.install(args)
|
||||||
|
|
||||||
@@ -995,6 +1011,110 @@ class BootstrapRuntimeTest(unittest.TestCase):
|
|||||||
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
|
self.assertIn(f"depth-anything-3 @ {da3_local}", ai_command)
|
||||||
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
|
self.assertIn(f"ben2 @ {ben2_local}", ai_command)
|
||||||
|
|
||||||
|
def test_install_stages_download_models_into_prefixed_directories(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
root = Path(directory)
|
||||||
|
runtime_dir = root / "runtime"
|
||||||
|
runtime_dir.mkdir()
|
||||||
|
sources = root / "sources"
|
||||||
|
sources.mkdir()
|
||||||
|
(sources / "depth-anything-3.tar.gz").write_bytes(b"da3")
|
||||||
|
(sources / "ben2.tar.gz").write_bytes(b"ben2")
|
||||||
|
uv_path = root / "uv"
|
||||||
|
uv_path.write_bytes(b"uv")
|
||||||
|
args = SimpleNamespace(
|
||||||
|
uv=str(uv_path),
|
||||||
|
runtime_dir=str(runtime_dir),
|
||||||
|
models_dir=str(root / "models"),
|
||||||
|
backend="cu128",
|
||||||
|
da3_commit="test-commit",
|
||||||
|
ben2_commit="test-ben2-commit",
|
||||||
|
status=str(root / "status.json"),
|
||||||
|
)
|
||||||
|
|
||||||
|
stages = bootstrap.install_stages(
|
||||||
|
args,
|
||||||
|
runtime_dir,
|
||||||
|
root / "python",
|
||||||
|
uv_path,
|
||||||
|
)
|
||||||
|
|
||||||
|
downloads = {
|
||||||
|
label: command
|
||||||
|
for label, command, _environment, kind in stages
|
||||||
|
if kind == "download"
|
||||||
|
}
|
||||||
|
self.assertIn("da3_small", " ".join(map(str, downloads["da3 small"])))
|
||||||
|
self.assertIn("da3_base", " ".join(map(str, downloads["da3 base"])))
|
||||||
|
self.assertIn("ben2_base", " ".join(map(str, downloads["ben2 base"])))
|
||||||
|
|
||||||
|
def test_run_verify_returns_detected_devices(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
|
||||||
|
class FakeProcess:
|
||||||
|
def __init__(self, lines):
|
||||||
|
self.stdout = lines
|
||||||
|
|
||||||
|
def wait(self):
|
||||||
|
return 0
|
||||||
|
|
||||||
|
fake_process = FakeProcess(
|
||||||
|
['SNAPFORM_DEVICES={"cuda": false, "mps": true}']
|
||||||
|
)
|
||||||
|
with patch.object(
|
||||||
|
bootstrap.subprocess,
|
||||||
|
"Popen",
|
||||||
|
return_value=fake_process,
|
||||||
|
):
|
||||||
|
report = bootstrap.run_verify(
|
||||||
|
"python",
|
||||||
|
"status.json",
|
||||||
|
clock=lambda: 0.0,
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertFalse(report["cuda"])
|
||||||
|
self.assertTrue(report["mps"])
|
||||||
|
|
||||||
|
def test_run_verify_writes_heartbeat_progress(self):
|
||||||
|
bootstrap = load_bootstrap_runtime()
|
||||||
|
|
||||||
|
class FakeProcess:
|
||||||
|
def __init__(self, lines):
|
||||||
|
self.stdout = lines
|
||||||
|
|
||||||
|
def wait(self):
|
||||||
|
return 0
|
||||||
|
|
||||||
|
clock_state = {"now": 0.0}
|
||||||
|
|
||||||
|
def fake_clock():
|
||||||
|
clock_state["now"] += 6.0
|
||||||
|
return clock_state["now"]
|
||||||
|
|
||||||
|
def delayed_empty_lines():
|
||||||
|
time.sleep(0.6)
|
||||||
|
if False:
|
||||||
|
yield
|
||||||
|
|
||||||
|
fake_process = FakeProcess(delayed_empty_lines())
|
||||||
|
with patch.object(
|
||||||
|
bootstrap.subprocess,
|
||||||
|
"Popen",
|
||||||
|
return_value=fake_process,
|
||||||
|
), patch.object(bootstrap, "write_status") as mocked_write:
|
||||||
|
report = bootstrap.run_verify(
|
||||||
|
"python",
|
||||||
|
"status.json",
|
||||||
|
clock=fake_clock,
|
||||||
|
)
|
||||||
|
|
||||||
|
messages = [call.args[2] for call in mocked_write.call_args_list]
|
||||||
|
self.assertTrue(
|
||||||
|
any("verifying runtime" in message for message in messages)
|
||||||
|
)
|
||||||
|
self.assertEqual(report, {})
|
||||||
|
|
||||||
def test_install_uv_extracts_bundled_archive(self):
|
def test_install_uv_extracts_bundled_archive(self):
|
||||||
bootstrap = load_bootstrap_runtime()
|
bootstrap = load_bootstrap_runtime()
|
||||||
with tempfile.TemporaryDirectory() as directory:
|
with tempfile.TemporaryDirectory() as directory:
|
||||||
|
|||||||
@@ -209,6 +209,21 @@ class ImageDepthTest(unittest.TestCase):
|
|||||||
(1, 1),
|
(1, 1),
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def test_extract_frame_range_honors_single_frame_setting(self):
|
||||||
|
source_object = SimpleNamespace(
|
||||||
|
data=SimpleNamespace(source="MOVIE", frame_duration=97),
|
||||||
|
image_user=SimpleNamespace(
|
||||||
|
frame_start=1,
|
||||||
|
frame_offset=70,
|
||||||
|
frame_duration=1,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
|
||||||
|
self.assertEqual(
|
||||||
|
self.image_depth.extract_frame_range(source_object),
|
||||||
|
(71, 1),
|
||||||
|
)
|
||||||
|
|
||||||
def test_configure_image_user_preserves_cyclic_setting(self):
|
def test_configure_image_user_preserves_cyclic_setting(self):
|
||||||
image_user = SimpleNamespace(
|
image_user = SimpleNamespace(
|
||||||
frame_start=0,
|
frame_start=0,
|
||||||
|
|||||||
@@ -54,3 +54,21 @@ class ModelTest(unittest.TestCase):
|
|||||||
("config.json", "model.safetensors"),
|
("config.json", "model.safetensors"),
|
||||||
)
|
)
|
||||||
self.assertNotIn("BEN2_BASE", self.models.MODELS)
|
self.assertNotIn("BEN2_BASE", self.models.MODELS)
|
||||||
|
|
||||||
|
def test_da3_model_directories_share_the_da3_prefix(self):
|
||||||
|
for model_key, model in self.models.MODELS.items():
|
||||||
|
with self.subTest(model_key=model_key):
|
||||||
|
self.assertTrue(
|
||||||
|
model.directory_name.startswith("da3_"),
|
||||||
|
model.directory_name,
|
||||||
|
)
|
||||||
|
self.assertEqual(
|
||||||
|
self.models.model_directory_name(model_key),
|
||||||
|
model.directory_name,
|
||||||
|
)
|
||||||
|
|
||||||
|
def test_ben2_model_directory_uses_ben2_prefix(self):
|
||||||
|
self.assertEqual(
|
||||||
|
self.models.model_directory_name("BEN2_BASE"),
|
||||||
|
"ben2_base",
|
||||||
|
)
|
||||||
|
|||||||
@@ -385,79 +385,6 @@ class WorkerTest(unittest.TestCase):
|
|||||||
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
with self.assertRaisesRegex(RuntimeError, "Unsupported input type"):
|
||||||
self.worker.prepare_input(input_path, 16, directory)
|
self.worker.prepare_input(input_path, 16, directory)
|
||||||
|
|
||||||
def test_download_passes_optional_file_patterns(self):
|
|
||||||
calls = []
|
|
||||||
fake_hugging_face = ModuleType("huggingface_hub")
|
|
||||||
fake_hugging_face.snapshot_download = lambda **options: calls.append(
|
|
||||||
options
|
|
||||||
)
|
|
||||||
args = SimpleNamespace(
|
|
||||||
repository="PramaLLC/BEN2",
|
|
||||||
model_dir="model",
|
|
||||||
status="status.json",
|
|
||||||
allow_pattern=["config.json", "model.safetensors"],
|
|
||||||
)
|
|
||||||
|
|
||||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
||||||
sys.modules,
|
|
||||||
{"huggingface_hub": fake_hugging_face},
|
|
||||||
), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"create_download_progress_class",
|
|
||||||
return_value=object,
|
|
||||||
):
|
|
||||||
args.model_dir = str(Path(directory) / "model")
|
|
||||||
args.status = str(Path(directory) / "status.json")
|
|
||||||
self.worker.download_model(args)
|
|
||||||
|
|
||||||
self.assertEqual(
|
|
||||||
calls[0]["allow_patterns"],
|
|
||||||
["config.json", "model.safetensors"],
|
|
||||||
)
|
|
||||||
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
|
|
||||||
|
|
||||||
def test_download_model_falls_back_to_modelscope_when_direct_fails(self):
|
|
||||||
calls = []
|
|
||||||
fake_hugging_face = ModuleType("huggingface_hub")
|
|
||||||
|
|
||||||
def fake_snapshot_download(**options):
|
|
||||||
calls.append(options)
|
|
||||||
raise OSError("Connection reset by peer")
|
|
||||||
|
|
||||||
fake_hugging_face.snapshot_download = fake_snapshot_download
|
|
||||||
args = SimpleNamespace(
|
|
||||||
repository="PramaLLC/BEN2",
|
|
||||||
model_dir="model",
|
|
||||||
status="status.json",
|
|
||||||
allow_pattern=[],
|
|
||||||
)
|
|
||||||
|
|
||||||
with tempfile.TemporaryDirectory() as directory, patch.dict(
|
|
||||||
sys.modules,
|
|
||||||
{"huggingface_hub": fake_hugging_face},
|
|
||||||
), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"create_download_progress_class",
|
|
||||||
return_value=object,
|
|
||||||
), patch.dict(os.environ, {"HF_ENDPOINT": ""}), patch.object(
|
|
||||||
self.worker,
|
|
||||||
"download_from_modelscope",
|
|
||||||
) as mocked_modelscope:
|
|
||||||
def fake_download_from_modelscope(*args, **kwargs):
|
|
||||||
model_dir = Path(args[1])
|
|
||||||
model_dir.mkdir(parents=True, exist_ok=True)
|
|
||||||
(model_dir / "config.json").write_text("{}")
|
|
||||||
(model_dir / "model.safetensors").write_bytes(b"x")
|
|
||||||
|
|
||||||
mocked_modelscope.side_effect = fake_download_from_modelscope
|
|
||||||
args.model_dir = str(Path(directory) / "model")
|
|
||||||
args.status = str(Path(directory) / "status.json")
|
|
||||||
self.worker.download_model(args)
|
|
||||||
|
|
||||||
self.assertEqual(len(calls), 1)
|
|
||||||
self.assertEqual(calls[0]["endpoint"], self.worker.HF_ENDPOINT)
|
|
||||||
self.assertEqual(mocked_modelscope.call_count, 1)
|
|
||||||
|
|
||||||
def test_download_model_skips_unreachable_endpoint(self):
|
def test_download_model_skips_unreachable_endpoint(self):
|
||||||
calls = []
|
calls = []
|
||||||
fake_hugging_face = ModuleType("huggingface_hub")
|
fake_hugging_face = ModuleType("huggingface_hub")
|
||||||
|
|||||||
71
uv.lock
generated
71
uv.lock
generated
@@ -2,7 +2,78 @@ version = 1
|
|||||||
revision = 3
|
revision = 3
|
||||||
requires-python = ">=3.11"
|
requires-python = ">=3.11"
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "colorama"
|
||||||
|
version = "0.4.6"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697, upload-time = "2022-10-25T02:36:22.414Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "iniconfig"
|
||||||
|
version = "2.3.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "packaging"
|
||||||
|
version = "26.3"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pluggy"
|
||||||
|
version = "1.6.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pygments"
|
||||||
|
version = "2.20.0"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/c3/b2/bc9c9196916376152d655522fdcebac55e66de6603a76a02bca1b6414f6c/pygments-2.20.0.tar.gz", hash = "sha256:6757cd03768053ff99f3039c1a36d6c0aa0b263438fcab17520b30a303a82b5f", size = 4955991, upload-time = "2026-03-29T13:29:33.898Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/f4/7e/a72dd26f3b0f4f2bf1dd8923c85f7ceb43172af56d63c7383eb62b332364/pygments-2.20.0-py3-none-any.whl", hash = "sha256:81a9e26dd42fd28a23a2d169d86d7ac03b46e2f8b59ed4698fb4785f946d0176", size = 1231151, upload-time = "2026-03-29T13:29:30.038Z" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[[package]]
|
||||||
|
name = "pytest"
|
||||||
|
version = "9.1.1"
|
||||||
|
source = { registry = "https://pypi.org/simple" }
|
||||||
|
dependencies = [
|
||||||
|
{ name = "colorama", marker = "sys_platform == 'win32'" },
|
||||||
|
{ name = "iniconfig" },
|
||||||
|
{ name = "packaging" },
|
||||||
|
{ name = "pluggy" },
|
||||||
|
{ name = "pygments" },
|
||||||
|
]
|
||||||
|
sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" }
|
||||||
|
wheels = [
|
||||||
|
{ url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" },
|
||||||
|
]
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "snapform"
|
name = "snapform"
|
||||||
version = "0.10.10"
|
version = "0.10.10"
|
||||||
source = { virtual = "." }
|
source = { virtual = "." }
|
||||||
|
|
||||||
|
[package.dev-dependencies]
|
||||||
|
dev = [
|
||||||
|
{ name = "pytest" },
|
||||||
|
]
|
||||||
|
|
||||||
|
[package.metadata]
|
||||||
|
|
||||||
|
[package.metadata.requires-dev]
|
||||||
|
dev = [{ name = "pytest" }]
|
||||||
|
|||||||
Reference in New Issue
Block a user