icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance device management and model handling in SnapForm

Browse Source 
- 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:
马南
2026-08-08 20:16:11 +08:00
parent aee460fe7f
commit 9c896213a5
18 changed files with 699 additions and 196 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/snapform/ai_host.py
View File

@@ -147,7 +147,7 @@ class ModelCache:
Path(self.ben2_key[0]).name if self.ben2_key is not None else ""
),
"da3": (
f"da3_{Path(self.da3_key[0]).name}"
Path(self.da3_key[0]).name
if self.da3_key is not None
else ""
),

83
src/snapform/bootstrap_runtime.py
View File

@@ -50,11 +50,15 @@ PLAIN_SLICE_SPAN = 0.02
RAMP_RATE = 0.02
CREEP_CAP = 0.60
CHECK_CODE = (
"import torch; "
"import json; import torch; "
"from depth_anything_3.api import DepthAnything3; "
"from ben2 import BEN_Base; "
"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",
model.repository,
"--model-dir",
models_dir / model.key.lower(),
models_dir / model.directory_name,
"--status",
args.status,
]
@@ -405,6 +409,70 @@ def verify_command(python):
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():
machine = platform.machine().lower()
if sys.platform == "win32":
@@ -643,9 +711,15 @@ def install(args):
installed_packages.extend(stage_installed)
runtime_offset = stage_end
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")
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)
ready_path = runtime_dir / "snapform_ready.json"
ready_path.write_text(
@@ -658,6 +732,7 @@ def install(args):
"fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION,
"backend": args.backend,
"devices": devices,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
"packages": sorted(set(installed_packages)),

17
src/snapform/image_depth.py
View File

@@ -97,7 +97,7 @@ def extract_frame_range(source_object):
)
frame_count = (
min(image_frames, total_frames)
if image_frames > 1
if image_frames >= 1
else total_frames
)
start_frame = min(1 + frame_offset, total_frames - frame_count + 1)
@@ -1400,12 +1400,14 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy
from .runtime import host_busy, model_ready
from .models import BEN2_MODEL
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and model_ready(BEN2_MODEL.key)
)
def execute(self, context):
@@ -1511,12 +1513,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy
from .runtime import host_busy, model_ready
from .models import MODELS
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and any(model_ready(model_key) for model_key in MODELS)
)
def invoke(self, context, _event):
@@ -1535,6 +1539,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return context.window_manager.invoke_props_dialog(self, width=340)
def draw(self, context):
from .models import BEN2_MODEL
from .runtime import model_ready
layout = self.layout
layout.prop(self, "snapform_depth_source")
if self.snapform_depth_source == "ESTIMATE":
@@ -1563,7 +1570,9 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
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")
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")
layout.prop(self, "snapform_delete_original")

12
src/snapform/models.py
View File

@@ -5,6 +5,7 @@ from dataclasses import dataclass
class ModelSpec:
key: str
label: str
directory_name: str
repository: str
license_name: str
description: str
@@ -16,6 +17,7 @@ MODELS = {
"SMALL": ModelSpec(
key="SMALL",
label="DA3 Small",
directory_name="da3_small",
repository="depth-anything/DA3-SMALL",
license_name="Apache 2.0",
description="Fastest any-view model with Apache 2.0 weights",
@@ -24,6 +26,7 @@ MODELS = {
"BASE": ModelSpec(
key="BASE",
label="DA3 Base",
directory_name="da3_base",
repository="depth-anything/DA3-BASE",
license_name="Apache 2.0",
description="Balanced any-view model with Apache 2.0 weights",
@@ -32,6 +35,7 @@ MODELS = {
"LARGE_1_1": ModelSpec(
key="LARGE_1_1",
label="DA3 Large 1.1",
directory_name="da3_large_1_1",
repository="depth-anything/DA3-LARGE-1.1",
license_name="CC BY-NC 4.0",
description="Higher quality non-commercial model with larger VRAM use",
@@ -40,6 +44,7 @@ MODELS = {
"GIANT_1_1": ModelSpec(
key="GIANT_1_1",
label="DA3 Giant 1.1",
directory_name="da3_giant_1_1",
repository="depth-anything/DA3-GIANT-1.1",
license_name="CC BY-NC 4.0",
description="Non-commercial Giant model with Gaussian head",
@@ -48,6 +53,7 @@ MODELS = {
"NESTED_GIANT_LARGE_1_1": ModelSpec(
key="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",
license_name="CC BY-NC 4.0",
description="Non-commercial nested model with metric depth and Gaussian head",
@@ -58,6 +64,7 @@ MODELS = {
BEN2_MODEL = ModelSpec(
key="BEN2_BASE",
label="BEN2 Base",
directory_name="ben2_base",
repository="PramaLLC/BEN2",
license_name="MIT",
description="Fast high-quality background removal model",
@@ -142,3 +149,8 @@ def get_downloadable_model(model_key):
return DOWNLOADABLE_MODELS[model_key]
except KeyError as 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

12
src/snapform/operators.py
View File

@@ -337,6 +337,7 @@ class RunJob(bpy.types.Operator):
paths["result"],
BEN2_MODEL.key,
start_frame=start_frame,
device=self._task_device(settings),
),
paths,
)
@@ -398,10 +399,21 @@ class RunJob(bpy.types.Operator):
self.job_type == IMAGE_DEPTH_JOB
and remove_background
),
device=self._task_device(settings),
),
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):
if not bpy.app.online_access:
raise RuntimeError(

37
src/snapform/preferences.py
View File

@@ -1,9 +1,9 @@
import shutil
import sys
from pathlib import Path
import bpy
from .runtime import device_enum_items
class SnapFormPreferences(bpy.types.AddonPreferences):
bl_idname = __package__
@@ -17,31 +17,14 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
subtype="DIR_PATH",
default=str(Path.home() / ".snapform"),
)
uv_executable: bpy.props.StringProperty(
name="uv Executable",
device: bpy.props.EnumProperty(
name="Device",
description=(
"Path to uv; leave empty to auto-download a standalone uv "
"when it is missing"
"Inference device used by production tasks; only devices supported "
"on this machine are listed"
),
subtype="FILE_PATH",
default=shutil.which("uv") or "",
)
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="",
items=device_enum_items,
default=0,
)
def draw(self, _context):
@@ -50,9 +33,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
layout = self.layout
layout.prop(self, "storage_root")
layout.prop(self, "uv_executable")
layout.prop(self, "torch_backend")
layout.prop(self, "pypi_index_url")
layout.prop(self, "device")
layout.separator()
layout.label(text=f"Runtime: {runtime_directory()}")
if runtime_ready():

17
src/snapform/properties.py
View File

@@ -1,6 +1,7 @@
import bpy
from .models import gs_model_enum_items, model_enum_items
from .runtime import device_enum_items
class SnapFormSettings(bpy.types.PropertyGroup):
@@ -18,7 +19,7 @@ class SnapFormSettings(bpy.types.PropertyGroup):
name="Extract Frames",
description="Maximum number of consecutive frames extracted from the start of a video or frame sequence",
default=16,
min=2,
min=1,
soft_max=64,
)
model: bpy.props.EnumProperty(
@@ -32,15 +33,13 @@ class SnapFormSettings(bpy.types.PropertyGroup):
items=gs_model_enum_items,
)
device: bpy.props.EnumProperty(
name="Device",
description="Inference device used by the external Runtime",
items=(
("AUTO", "Auto", "Use CUDA, then MPS, otherwise CPU"),
("CUDA", "CUDA", "Require an NVIDIA CUDA device"),
("MPS", "MPS", "Use Apple Metal (MPS)"),
("CPU", "CPU", "Run on CPU"),
name="Debug Device",
description=(
"Inference device used by Debug tasks only; production tasks use "
"the Device setting in the Add-on Preferences"
),
default="AUTO",
items=device_enum_items,
default=0,
)
process_res: bpy.props.IntProperty(
name="Process Res",

96
src/snapform/runtime.py
View File

@@ -1,5 +1,6 @@
import json
import os
import platform
import secrets
import shutil
import socket
@@ -15,7 +16,12 @@ from urllib.request import ProxyHandler, Request, build_opener
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"
@@ -107,7 +113,16 @@ def models_directory():
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):
@@ -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():
path = storage_root() / "jobs"
path.mkdir(parents=True, exist_ok=True)
@@ -153,13 +223,7 @@ def new_job_paths(job_type):
def uv_executable():
addon_preferences = preferences()
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 the bundled uv path inside the shared storage."""
return str(uv_cache_path())
@@ -171,9 +235,7 @@ def uv_cache_path():
def install_runtime_command(status_path):
executable = uv_executable()
addon_preferences = preferences()
backend = getattr(addon_preferences, "torch_backend", "CU128")
pypi_index = getattr(addon_preferences, "pypi_index_url", "") or ""
backend = "cu128" if "cuda" in machine_devices() else "cpu"
script = Path(__file__).with_name("bootstrap_runtime.py")
command = [
sys.executable,
@@ -185,7 +247,7 @@ def install_runtime_command(status_path):
"--models-dir",
str(models_directory()),
"--backend",
backend.lower(),
backend,
"--da3-commit",
DA3_COMMIT,
"--ben2-commit",
@@ -193,8 +255,6 @@ def install_runtime_command(status_path):
"--status",
str(status_path),
]
if pypi_index:
command.extend(["--pypi-index", pypi_index])
return command
@@ -640,6 +700,7 @@ def remove_background_request(
model_key,
video_frames=None,
start_frame=1,
device=None,
):
return HostJobRequest(
{
@@ -650,7 +711,7 @@ def remove_background_request(
"model_dir": str(model_directory(model_key)),
"output_dir": str(result_directory),
"status": str(status_path),
"device": settings.device.lower(),
"device": (device or settings.device).lower(),
"video_frames": (
settings.video_frames
if video_frames is None
@@ -701,6 +762,7 @@ def reconstruct_request(
use_ray_pose=None,
remove_background=False,
start_frame=1,
device=None,
):
snapform_selected_key = model_key or selected_model_key(settings, export_format)
snapform_selected_process_res = (
@@ -732,7 +794,7 @@ def reconstruct_request(
"gs_views_interval": settings.gs_views_interval,
"video_frames": settings.video_frames,
"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,
"ref_view_strategy": settings.ref_view_strategy,
"remove_background_model_dir": "",