Add unit tests for service functionality and model caching

- Implement tests for ben2 and da3 model caching to ensure models are reused for matching runtime keys.
- Verify service state management, including client registration, job cancellation, and client expiration.
- Test discovery file creation and removal based on service ownership.
- Ensure frame extraction functionality works without models and produces expected results.
This commit is contained in:
2026-08-13 00:35:40 +08:00
parent 7d0983529e
commit 92d2cea1d1
21 changed files with 935 additions and 887 deletions

View File

@@ -1,6 +1,6 @@
import bpy
from . import runtime
from . import client
from .image_depth import (
ApplyRemovedBackground,
BuildDepthPlane,
@@ -13,7 +13,7 @@ from .image_depth import (
from .image_mesh import BuildImageMesh, ConvertImageToMesh
from .operators import (
CancelJob,
ControlHost,
ControlService,
RunJob,
cancel_active_job,
close_active_job,
@@ -22,7 +22,7 @@ from .panel import (
Ben2RemovalDebugPanel,
Da3OutputDebugPanel,
DebugPanel,
HostPanel,
AiServicePanel,
ModelsPanel,
)
from .preferences import SnapFormPreferences
@@ -35,7 +35,7 @@ CLASSES = (
SnapFormSettings,
RunJob,
CancelJob,
ControlHost,
ControlService,
SnapFormImageMenu,
ApplyRemovedBackground,
RemoveImageBackground,
@@ -44,7 +44,7 @@ CLASSES = (
ConvertImageToDepthPlane,
BuildImageMesh,
ConvertImageToMesh,
HostPanel,
AiServicePanel,
ModelsPanel,
DebugPanel,
Da3OutputDebugPanel,
@@ -61,14 +61,14 @@ def register():
bpy.types.STATUSBAR_HT_header.append(draw_status_bar)
bpy.types.VIEW3D_MT_object_context_menu.append(draw_image_context_menu)
bpy.types.OUTLINER_MT_object.append(draw_image_context_menu)
runtime.schedule_host_start()
client.schedule_service_start()
def unregister():
cancel_active_job()
close_active_job()
runtime.unschedule_host()
runtime.detach_host_client()
client.unschedule_service()
client.detach_service_client()
bpy.types.VIEW3D_MT_object_context_menu.remove(draw_image_context_menu)
bpy.types.OUTLINER_MT_object.remove(draw_image_context_menu)
bpy.types.STATUSBAR_HT_header.remove(draw_status_bar)

View File

@@ -12,8 +12,8 @@ platforms = ["windows-x64", "macos-arm64", "linux-x64"]
license = ["SPDX:GPL-3.0-or-later"]
[permissions]
network = "Access the local Backend Service and download environment dependencies and model weights"
files = "Read source images and store runtime data"
network = "Access the local Service and download environment dependencies and model weights"
files = "Read source images and store environment data"
[build]
paths_exclude_pattern = [

View File

@@ -28,13 +28,11 @@ from .models import (
DA3_COMMIT = "3d835ec1a5802d64a8b8b15f817a1ab54809bfe4"
BEN2_COMMIT = "2c99a5da477b5523585bfa5c893888a6e818a8f6"
ENVIRONMENT_MANIFEST_REVISION = 4
JOB_PROTOCOL_VERSION = 2
ENVIRONMENT_MANIFEST_FILENAME = "environment_manifest.json"
HOST_ADDRESS = "127.0.0.1"
MANIFEST_FILENAME = "manifest.json"
SERVICE_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
JOB_FOLDER_NAMES = {
"INSTALL_RUNTIME": "install-runtime",
"INSTALL_ENV": "install-environment",
"DOWNLOAD_MODEL": "download-model",
"RECONSTRUCT": "export-reconstruction",
"IMAGE_DEPTH": "generate-depth-plane",
@@ -47,12 +45,12 @@ JOB_FOLDER_NAMES = {
@dataclass(frozen=True)
class HostJobRequest:
class JobRequest:
payload: dict
@dataclass
class HostConnection:
class Connection:
port: int
token: str
pid: int
@@ -60,17 +58,17 @@ class HostConnection:
@dataclass
class LaunchedHost:
class ServiceProcess:
process: subprocess.Popen
log_file: object
_host_connection = None
_launched_host = None
_service_connection = None
_launched_service = None
_registered_instance_id = None
_client_id = uuid.uuid4().hex
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_service_auto_start = False
_service_snapshot = {"state": "STOPPED", "message": "Stopped"}
_local_opener = build_opener(ProxyHandler({}))
_client_label = f"Blender {getattr(bpy.app, 'version_string', '')}".strip()
@@ -92,29 +90,26 @@ def storage_root():
return root
def runtime_directory():
return storage_root() / "runtime"
def venv_directory():
return storage_root() / ".venv"
def runtime_python():
directory = runtime_directory()
def venv_python():
directory = venv_directory()
if os.name == "nt":
return directory / "Scripts" / "python.exe"
return directory / "bin" / "python"
def runtime_ready():
ready_path = runtime_directory() / ENVIRONMENT_MANIFEST_FILENAME
if not runtime_python().is_file() or not ready_path.is_file():
def env_installed():
ready_path = storage_root() / MANIFEST_FILENAME
if not venv_python().is_file() or not ready_path.is_file():
return False
try:
ready = json.loads(ready_path.read_text(encoding="utf-8"))
manifest = json.loads(ready_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
return False
return (
ready.get("environment_manifest_revision")
== ENVIRONMENT_MANIFEST_REVISION
)
return isinstance(manifest, dict)
def models_directory():
@@ -143,9 +138,9 @@ def model_ready(model_key):
)
def runtime_devices():
"""Devices detected by the Backend Environment, or None when unknown."""
ready_path = runtime_directory() / ENVIRONMENT_MANIFEST_FILENAME
def env_devices():
"""Devices detected by the installed environment, or None when unknown."""
ready_path = storage_root() / MANIFEST_FILENAME
if not ready_path.is_file():
return None
try:
@@ -159,7 +154,7 @@ def runtime_devices():
def machine_devices():
"""Devices the machine can run before Backend Environment installation."""
"""Devices the machine can run before environment installation."""
devices = ["cpu"]
if shutil.which("nvidia-smi"):
devices.insert(0, "cuda")
@@ -173,7 +168,7 @@ def machine_devices():
def supported_devices():
"""Inference devices available on this machine, fastest first."""
return runtime_devices() or machine_devices()
return env_devices() or machine_devices()
DEVICE_LABELS = {
@@ -249,12 +244,16 @@ def new_job_paths(job_type):
directory = jobs_directory() / stem
directory.mkdir(parents=True, exist_ok=True)
return {
"log": directory / f"{directory.name}.log",
"status": directory / f"{directory.name}.json",
"log": directory / "job.log",
"status": directory / "status.json",
"result": directory,
}
def job_id_from_status(status_path):
return Path(status_path).parent.name
def uv_executable():
"""Return the bundled uv path inside the shared storage."""
return str(uv_cache_path())
@@ -262,25 +261,27 @@ def uv_executable():
def uv_cache_path():
if os.name == "nt":
return storage_root() / "uv" / "uv.exe"
return storage_root() / "uv" / "uv"
return storage_root() / "tools" / "uv.exe"
return storage_root() / "tools" / "uv"
def install_runtime_command(status_path):
def install_env_command(status_path):
executable = uv_executable()
backend = "cu128" if "cuda" in machine_devices() else "cpu"
script = Path(__file__).with_name("bootstrap_runtime.py")
torch_variant = "cu128" if "cuda" in machine_devices() else "cpu"
script = Path(__file__).with_name("install_env.py")
command = [
sys.executable,
str(script),
"--uv",
executable,
"--runtime-dir",
str(runtime_directory()),
"--venv-dir",
str(venv_directory()),
"--manifest",
str(storage_root() / MANIFEST_FILENAME),
"--models-dir",
str(models_directory()),
"--backend",
backend,
"--torch-variant",
torch_variant,
"--da3-commit",
DA3_COMMIT,
"--ben2-commit",
@@ -291,36 +292,36 @@ def install_runtime_command(status_path):
return command
def host_log_path():
return storage_root() / "ai-host.log"
def service_log_path():
return storage_root() / "service.log"
def host_discovery_path():
return storage_root() / "host.json"
def service_discovery_path():
return storage_root() / "service.json"
def host_lock_path():
return storage_root() / "host.lock"
def service_lock_path():
return storage_root() / "service.lock"
def _available_host_port():
def _available_service_port():
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as listener:
listener.bind((HOST_ADDRESS, 0))
listener.bind((SERVICE_ADDRESS, 0))
return listener.getsockname()[1]
def _host_url(connection, path):
return f"http://{HOST_ADDRESS}:{connection.port}{path}"
def _service_url(connection, path):
return f"http://{SERVICE_ADDRESS}:{connection.port}{path}"
def _host_http_request(connection, method, path, payload=None, timeout=2.0):
def _service_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),
_service_url(connection, path),
data=data,
headers=headers,
method=method,
@@ -334,16 +335,16 @@ def _host_http_request(connection, method, path, payload=None, timeout=2.0):
except (json.JSONDecodeError, UnicodeDecodeError):
detail = None
raise RuntimeError(
detail or f"Backend Service returned HTTP {error.code}"
detail or f"Service returned HTTP {error.code}"
)
except (OSError, URLError) as error:
raise RuntimeError(f"Unable to reach Backend Service: {error}")
raise RuntimeError(f"Unable to reach Service: {error}")
def _read_host_discovery():
def _read_service_discovery():
try:
payload = json.loads(host_discovery_path().read_text(encoding="utf-8"))
return HostConnection(
payload = json.loads(service_discovery_path().read_text(encoding="utf-8"))
return Connection(
port=int(payload["port"]),
token=str(payload["token"]),
pid=int(payload["pid"]),
@@ -353,46 +354,44 @@ def _read_host_discovery():
return None
def _check_host(connection, timeout=0.3):
health = _host_http_request(
def _check_service(connection, timeout=0.3):
health = _service_http_request(
connection,
"GET",
"/health",
timeout=timeout,
)
if health.get("job_protocol_version") != JOB_PROTOCOL_VERSION:
raise RuntimeError("Backend Service job protocol is out of date")
if health.get("instance_id") != connection.instance_id:
raise RuntimeError("Backend Service discovery is stale")
raise RuntimeError("Service discovery is stale")
return health
def _discover_live_host(timeout=0.3):
connection = _read_host_discovery()
def _discover_live_service(timeout=0.3):
connection = _read_service_discovery()
if connection is None:
return None, None
try:
return connection, _check_host(connection, timeout=timeout)
return connection, _check_service(connection, timeout=timeout)
except RuntimeError:
return None, None
def start_host():
global _launched_host
def start_service():
global _launched_service
if _launched_host is not None and _launched_host.process.poll() is None:
if _launched_service is not None and _launched_service.process.poll() is None:
return
_close_launched_host()
port = _available_host_port()
_close_launched_service()
port = _available_service_port()
token = secrets.token_urlsafe(32)
instance_id = uuid.uuid4().hex
log_file = host_log_path().open("a", encoding="utf-8")
log_file = service_log_path().open("a", encoding="utf-8")
command = [
str(runtime_python()),
str(venv_python()),
"-u",
str(Path(__file__).with_name("ai_host.py")),
str(Path(__file__).with_name("service.py")),
"--host",
HOST_ADDRESS,
SERVICE_ADDRESS,
"--port",
str(port),
"--token",
@@ -400,9 +399,9 @@ def start_host():
"--instance-id",
instance_id,
"--discovery",
str(host_discovery_path()),
str(service_discovery_path()),
"--lock",
str(host_lock_path()),
str(service_lock_path()),
]
try:
process = subprocess.Popen(
@@ -415,46 +414,46 @@ def start_host():
except OSError:
log_file.close()
raise
_launched_host = LaunchedHost(process, log_file)
_launched_service = ServiceProcess(process, log_file)
def ensure_host(timeout=60.0):
global _host_auto_start, _host_connection, _host_snapshot
def ensure_service(timeout=60.0):
global _service_auto_start, _service_connection, _service_snapshot
_host_auto_start = True
_service_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)
connection, health = _discover_live_service(timeout=0.5)
if connection is not None:
_host_connection = connection
_register_host_client(connection)
_host_snapshot = _format_host_snapshot(connection, health)
_redraw_host_ui()
_service_connection = connection
_register_service_client(connection)
_service_snapshot = _format_service_snapshot(connection, health)
_redraw_service_ui()
return connection
if not starting:
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
_service_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_service_ui()
starting = True
_close_finished_launched_host()
_close_finished_launched_service()
if (
_launched_host is None
_launched_service is None
and time.monotonic() - last_launch >= 1.0
):
start_host()
start_service()
last_launch = time.monotonic()
time.sleep(0.05)
raise RuntimeError(
f"Backend Service did not become ready; log: {host_log_path()}"
f"Service did not become ready; log: {service_log_path()}"
)
def submit_host_job(request):
connection = _host_connection or ensure_host()
def submit_service_job(request):
connection = _service_connection or ensure_service()
payload = dict(request.payload)
payload["client_id"] = _client_id
_host_http_request(
_service_http_request(
connection,
"POST",
"/jobs",
@@ -464,79 +463,79 @@ def submit_host_job(request):
return connection
def stop_host():
global _host_auto_start, _host_connection, _host_snapshot
def stop_service():
global _service_auto_start, _service_connection, _service_snapshot
global _registered_instance_id
connection = _host_connection
connection = _service_connection
if connection is None:
connection, _health = _discover_live_host(timeout=0.5)
connection, _health = _discover_live_service(timeout=0.5)
if connection is None:
_close_launched_host(terminate=True)
_host_auto_start = False
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_close_launched_service(terminate=True)
_service_auto_start = False
_service_snapshot = {"state": "STOPPED", "message": "Stopped"}
return
_register_host_client(connection)
_host_http_request(
_register_service_client(connection)
_service_http_request(
connection,
"POST",
"/shutdown",
payload={"client_id": _client_id},
timeout=2.0,
)
_host_auto_start = False
_service_auto_start = False
deadline = time.monotonic() + 10.0
while time.monotonic() < deadline:
try:
_check_host(connection, timeout=0.2)
_check_service(connection, timeout=0.2)
except RuntimeError:
break
time.sleep(0.05)
else:
raise RuntimeError("Backend Service did not stop")
_host_connection = None
raise RuntimeError("Service did not stop")
_service_connection = None
_registered_instance_id = None
_close_finished_launched_host()
_host_snapshot = {"state": "STOPPED", "message": "Stopped"}
_close_finished_launched_service()
_service_snapshot = {"state": "STOPPED", "message": "Stopped"}
def restart_host():
global _host_auto_start, _host_connection, _registered_instance_id
global _host_snapshot
def restart_service():
global _service_auto_start, _service_connection, _registered_instance_id
global _service_snapshot
connection = ensure_host()
_host_http_request(
connection = ensure_service()
_service_http_request(
connection,
"POST",
"/shutdown",
payload={"client_id": _client_id},
timeout=2.0,
)
_host_auto_start = True
_host_connection = None
_service_auto_start = True
_service_connection = None
_registered_instance_id = None
_host_snapshot = {"state": "STARTING", "message": "Restarting"}
_service_snapshot = {"state": "STARTING", "message": "Restarting"}
def enable_host():
global _host_auto_start, _host_connection, _host_snapshot
def enable_service():
global _service_auto_start, _service_connection, _service_snapshot
_host_auto_start = True
connection, health = _discover_live_host(timeout=0.2)
_service_auto_start = True
connection, health = _discover_live_service(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()
_service_connection = connection
_register_service_client(connection)
_service_snapshot = _format_service_snapshot(connection, health)
_redraw_service_ui()
return
start_host()
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
start_service()
_service_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_service_ui()
def clear_host_models():
connection = ensure_host()
return _host_http_request(
def clear_service_models():
connection = ensure_service()
return _service_http_request(
connection,
"POST",
"/models/clear",
@@ -544,9 +543,9 @@ def clear_host_models():
)
def cancel_host_job(job_id):
connection = ensure_host()
return _host_http_request(
def cancel_service_job(job_id):
connection = ensure_service()
return _service_http_request(
connection,
"DELETE",
f"/jobs/{job_id}",
@@ -555,13 +554,13 @@ def cancel_host_job(job_id):
)
def detach_host_client():
global _host_connection, _registered_instance_id
def detach_service_client():
global _service_connection, _registered_instance_id
connection = _host_connection
connection = _service_connection
if connection is not None and _registered_instance_id is not None:
try:
_host_http_request(
_service_http_request(
connection,
"DELETE",
f"/clients/{_client_id}",
@@ -569,48 +568,48 @@ def detach_host_client():
)
except RuntimeError:
pass
_host_connection = None
_service_connection = None
_registered_instance_id = None
# 当前 Blender 启动的 Service 随 Blender 关闭,不等待空闲超时。
_close_launched_host(terminate=True)
_close_launched_service(terminate=True)
def schedule_host_start():
def schedule_service_start():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_host_timer):
if timers.is_registered(_service_timer):
return
timers.register(_service_host_timer, first_interval=1.0, persistent=True)
timers.register(_service_timer, first_interval=1.0, persistent=True)
def unschedule_host():
def unschedule_service():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_host_timer):
timers.unregister(_service_host_timer)
if timers.is_registered(_service_timer):
timers.unregister(_service_timer)
def host_status():
return dict(_host_snapshot)
def service_status():
return dict(_service_snapshot)
def host_busy():
return _host_snapshot.get("state") == "BUSY"
def service_busy():
return _service_snapshot.get("state") == "BUSY"
def mark_host_job_complete(job_id):
def mark_service_job_complete(job_id):
"""Release a stale local busy snapshot after its terminal status is observed."""
global _host_snapshot
global _service_snapshot
active = _host_snapshot.get("active_job")
active = _service_snapshot.get("active_job")
if active and active.get("job_id") != job_id:
return
if _host_snapshot.get("state") != "BUSY":
if _service_snapshot.get("state") != "BUSY":
return
_host_snapshot = {
**_host_snapshot,
_service_snapshot = {
**_service_snapshot,
"busy": False,
"active_job": None,
"state": "READY",
@@ -618,38 +617,38 @@ def mark_host_job_complete(job_id):
}
def _service_host_timer():
global _host_connection, _host_snapshot, _registered_instance_id
def _service_timer():
global _service_connection, _service_snapshot, _registered_instance_id
if not runtime_ready():
_host_snapshot = {
if not env_installed():
_service_snapshot = {
"state": "UNAVAILABLE",
"message": "Backend Environment is not installed",
"message": "Environment is not installed",
}
return 2.0
connection, health = _discover_live_host(timeout=0.2)
connection, health = _discover_live_service(timeout=0.2)
if connection is None:
_host_connection = None
_service_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()
_close_finished_launched_service()
if not _service_auto_start:
if _service_snapshot.get("state") != "STOPPED":
_service_snapshot = {"state": "STOPPED", "message": "Stopped"}
_redraw_service_ui()
return 2.0
try:
start_host()
start_service()
except OSError as error:
_host_snapshot = {"state": "ERROR", "message": str(error)}
_service_snapshot = {"state": "ERROR", "message": str(error)}
return 2.0
_host_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_host_ui()
_service_snapshot = {"state": "STARTING", "message": "Starting"}
_redraw_service_ui()
return 1.0
try:
_host_connection = connection
_service_connection = connection
if _registered_instance_id == connection.instance_id:
try:
_host_http_request(
_service_http_request(
connection,
"POST",
f"/clients/{_client_id}/heartbeat",
@@ -657,23 +656,23 @@ def _service_host_timer():
)
except RuntimeError:
_registered_instance_id = None
_register_host_client(connection)
health = _check_host(connection, timeout=0.5)
_host_snapshot = _format_host_snapshot(connection, health)
_register_service_client(connection)
health = _check_service(connection, timeout=0.5)
_service_snapshot = _format_service_snapshot(connection, health)
except RuntimeError as error:
_host_connection = None
_service_connection = None
_registered_instance_id = None
_host_snapshot = {"state": "ERROR", "message": str(error)}
_redraw_host_ui()
_service_snapshot = {"state": "ERROR", "message": str(error)}
_redraw_service_ui()
return 5.0
def _register_host_client(connection):
def _register_service_client(connection):
global _registered_instance_id
if _registered_instance_id == connection.instance_id:
return
_host_http_request(
_service_http_request(
connection,
"POST",
"/clients",
@@ -683,7 +682,7 @@ def _register_host_client(connection):
_registered_instance_id = connection.instance_id
def _format_host_snapshot(connection, health):
def _format_service_snapshot(connection, health):
state = "BUSY" if health.get("busy") else "READY"
return {
**health,
@@ -693,7 +692,7 @@ def _format_host_snapshot(connection, health):
}
def _redraw_host_ui():
def _redraw_service_ui():
if threading.current_thread() is not threading.main_thread():
return
try:
@@ -704,17 +703,17 @@ def _redraw_host_ui():
pass
def _close_finished_launched_host():
if _launched_host is None or _launched_host.process.poll() is None:
def _close_finished_launched_service():
if _launched_service is None or _launched_service.process.poll() is None:
return
_close_launched_host()
_close_launched_service()
def _close_launched_host(terminate=False):
global _launched_host
def _close_launched_service(terminate=False):
global _launched_service
launched = _launched_host
_launched_host = None
launched = _launched_service
_launched_service = None
if launched is None:
return
if terminate and launched.process.poll() is None:
@@ -730,7 +729,7 @@ def _close_launched_host(terminate=False):
def download_model_command(model_key, status_path):
model = get_downloadable_model(model_key)
command = [
str(runtime_python()),
str(venv_python()),
str(Path(__file__).with_name("worker.py")),
"download",
"--repository",
@@ -757,9 +756,9 @@ def remove_background_request(
start_frame=1,
device=None,
):
return HostJobRequest(
return JobRequest(
{
"job_id": Path(status_path).stem,
"job_id": job_id_from_status(status_path),
"command": "remove-background",
"parameters": {
"input": str(input_path),
@@ -786,9 +785,9 @@ def extract_frames_request(
video_frames=None,
start_frame=1,
):
return HostJobRequest(
return JobRequest(
{
"job_id": Path(status_path).stem,
"job_id": job_id_from_status(status_path),
"command": "extract-frames",
"parameters": {
"input": str(input_path),
@@ -867,9 +866,9 @@ def reconstruct_request(
parameters["remove_background_model_dir"] = str(
model_directory(BEN2_MODEL.key)
)
return HostJobRequest(
return JobRequest(
{
"job_id": Path(status_path).stem,
"job_id": job_id_from_status(status_path),
"command": "reconstruct",
"parameters": parameters,
}

View File

@@ -49,7 +49,7 @@ VIDEO_NAME_SUFFIXES = (
def image_depth_model_items(owner, context):
from .models import model_enum_items
from .runtime import model_ready
from .client import model_ready
return [
item
@@ -223,7 +223,7 @@ def draw_process_resolution(layout, owner, image):
def save_packed_image(image, scene):
from .runtime import temporary_input_path
from .client import temporary_input_path
output_path = temporary_input_path()
image_settings = scene.render.image_settings
@@ -281,7 +281,7 @@ def prepare_image_input(image, scene):
return source_path, False
sequence_frames = collect_sequence_frames(image)
if sequence_frames:
from .runtime import temporary_input_directory
from .client import temporary_input_directory
temporary_directory = Path(temporary_input_directory())
try:
@@ -881,12 +881,14 @@ def material_node_group(depth_plane=False):
return existing
group = bpy.data.node_groups.new(group_name, "ShaderNodeTree")
interface = group.interface
for name, socket_type in (
input_sockets = [
("Color", "NodeSocketColor"),
("Alpha", "NodeSocketFloat"),
("Depth", "NodeSocketFloat"),
("Bump Scale", "NodeSocketFloat"),
):
]
if depth_plane:
input_sockets.insert(2, ("Depth", "NodeSocketFloat"))
for name, socket_type in input_sockets:
socket = interface.new_socket(
name=name,
in_out="INPUT",
@@ -934,7 +936,6 @@ def material_node_group(depth_plane=False):
displacement = None
processed_alpha = group_input.outputs["Alpha"]
processed_depth = group_input.outputs["Depth"]
bump_depth = group_input.outputs["Color"]
if depth_plane:
displacement = nodes.new("ShaderNodeDisplacement")
@@ -1174,8 +1175,6 @@ def create_material(
output.inputs["Displacement"],
)
set_displacement_only(material)
else:
links.new(color_texture.outputs["Color"], material_inputs.inputs["Depth"])
return material
@@ -1700,13 +1699,13 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .client import service_busy, model_ready
from .models import BEN2_MODEL
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and not service_busy()
and model_ready(BEN2_MODEL.key)
)
@@ -1790,13 +1789,13 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .client import service_busy, model_ready
from .models import MODELS
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and not service_busy()
and any(model_ready(model_key) for model_key in MODELS)
)
@@ -1823,7 +1822,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
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 AI Backend > Models")
estimate_box.label(text="Download one in Models")
else:
estimate_box.prop(self, "snapform_model")
estimate_box.prop(self, "snapform_use_ray_pose")
@@ -1917,17 +1916,17 @@ class SnapFormImageMenu(bpy.types.Menu):
bl_label = "SnapForm"
def draw(self, _context):
from .runtime import runtime_ready
from .client import env_installed
layout = self.layout
layout.operator_context = "INVOKE_DEFAULT"
if not runtime_ready():
if not env_installed():
operator = layout.operator(
"snapform.run_job",
text="Install Backend",
text="Install Environment",
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
operator.job_type = "INSTALL_ENV"
return
layout.operator(
"snapform.convert_image_to_depth_plane",

View File

@@ -517,13 +517,13 @@ class ConvertImageToMesh(bpy.types.Operator):
def poll(cls, context):
from .models import MODELS
from .operators import active_job
from .runtime import host_busy, model_ready
from .client import service_busy, model_ready
return (
is_image_empty(context.object)
and video_total_frames(context.object.data) == 1
and active_job() is None
and not host_busy()
and not service_busy()
and any(model_ready(model_key) for model_key in MODELS)
)

View File

@@ -34,8 +34,6 @@ XFORMERS_VERSION = "0.0.33.post2"
ADDICT_VERSION = "2.4.0"
FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0"
ENVIRONMENT_MANIFEST_REVISION = 4
ENVIRONMENT_MANIFEST_FILENAME = "environment_manifest.json"
TORCH_INDEXES = {
"cu128": "https://download.pytorch.org/whl/cu128",
"cpu": "https://download.pytorch.org/whl/cpu",
@@ -148,8 +146,8 @@ def estimated_remaining(label, progress, elapsed=0.0):
return max(expected_seconds * remaining_fraction - elapsed, 0.0)
def effective_torch_index(backend):
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(backend)
def effective_torch_index(torch_variant):
return os.environ.get("SNAPFORM_TORCH_INDEX") or TORCH_INDEXES.get(torch_variant)
def resolve_indexes(args):
@@ -304,21 +302,23 @@ def run_install(
print(f"Retrying with PyPI mirror: {PYPI_MIRRORS[0]}")
def environment_python(runtime_dir):
def venv_python(venv_dir):
if sys.platform == "win32":
return Path(runtime_dir) / "Scripts" / "python.exe"
return Path(runtime_dir) / "bin" / "python"
return Path(venv_dir) / "Scripts" / "python.exe"
return Path(venv_dir) / "bin" / "python"
def install_stages(args, runtime_dir, python, uv):
def install_stages(args, venv_dir, python, uv, vendor_dir=None):
# macOS 的 PyTorch wheel 只在 PyPI 上,download.pytorch.org 没有 macosx 构建
torch_index = (
None
if sys.platform == "darwin"
else effective_torch_index(args.backend)
else effective_torch_index(args.torch_variant)
)
vendor_dir = Path(vendor_dir or Path(venv_dir).parent / "vendor")
def archive_requirement(project_name, filename):
local = Path(args.runtime_dir).parent / "vendor" / filename
local = Path(vendor_dir) / filename
if not local.is_file():
raise RuntimeError(
"Bundled vendor artifact is missing; reinstall the extension: "
@@ -349,7 +349,7 @@ def install_stages(args, runtime_dir, python, uv):
stages = [
(
"Environment",
[uv, "venv", "--clear", "--python", "3.12", runtime_dir],
[uv, "venv", "--clear", "--python", "3.12", venv_dir],
None,
"plain",
),
@@ -390,7 +390,7 @@ def install_stages(args, runtime_dir, python, uv):
python,
f"xformers=={XFORMERS_VERSION}",
]
if args.backend == "cu128":
if args.torch_variant == "cu128":
xformers_command.extend(["--index-url", torch_index])
else:
xformers_command.append("--no-deps")
@@ -429,18 +429,15 @@ def install_stages(args, runtime_dir, python, uv):
return stages
def runtime_ready(runtime_dir):
ready_path = runtime_dir / ENVIRONMENT_MANIFEST_FILENAME
if not environment_python(runtime_dir).is_file() or not ready_path.is_file():
def env_installed(venv_dir, manifest_path=None):
ready_path = Path(manifest_path or Path(venv_dir).parent / "manifest.json")
if not venv_python(venv_dir).is_file() or not ready_path.is_file():
return False
try:
payload = json.loads(ready_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
return False
return (
payload.get("environment_manifest_revision")
== ENVIRONMENT_MANIFEST_REVISION
)
return isinstance(payload, dict)
def installed_devices(args):
@@ -448,7 +445,7 @@ def installed_devices(args):
machine = platform.machine().lower()
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
devices.append("mps")
elif args.backend.startswith("cu"):
elif args.torch_variant.startswith("cu"):
devices.append("cuda")
devices.append("cpu")
return devices
@@ -468,7 +465,7 @@ def run_verify(
stages=None,
stage_label=None,
):
"""Run the runtime import check with heartbeat progress; return the device report."""
"""Run the environment 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(
@@ -509,7 +506,7 @@ def run_verify(
write_status(
status_path,
progress,
"Checking Backend Environment",
"Checking Environment",
stage=stage,
stages=stages,
stage_label=stage_label,
@@ -560,8 +557,8 @@ def bundled_uv_archive():
return path if path.is_file() else None
def sync_bundled_vendor(bundled_dir, vendor_dir):
"""Copy bundled third-party archives into persistent storage."""
def copy_vendor_inputs(bundled_dir, vendor_dir):
"""Copy bundled installation inputs into a temporary directory."""
for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"):
bundled = Path(bundled_dir) / name
if not bundled.is_file():
@@ -573,6 +570,14 @@ def sync_bundled_vendor(bundled_dir, vendor_dir):
shutil.copy2(bundled, target)
def remove_legacy_vendor(storage_root):
"""Remove installation inputs persisted by older development builds."""
root = Path(storage_root)
shutil.rmtree(root / "vendor", ignore_errors=True)
for name in ("depth-anything-3.tar.gz", "ben2.tar.gz"):
(root / name).unlink(missing_ok=True)
def archive_is_valid(path):
if zipfile.is_zipfile(path):
return True
@@ -714,22 +719,28 @@ def install(args):
"shipping Intel Mac (x86_64) wheels after 2.2.2"
)
resolve_indexes(args)
runtime_dir = Path(args.runtime_dir)
runtime_dir.parent.mkdir(parents=True, exist_ok=True)
if not runtime_ready(runtime_dir):
shutil.rmtree(runtime_dir, ignore_errors=True)
venv_dir = Path(args.venv_dir)
manifest_path = Path(
getattr(args, "manifest", venv_dir.parent / "manifest.json")
)
venv_dir.parent.mkdir(parents=True, exist_ok=True)
if not env_installed(venv_dir, manifest_path):
shutil.rmtree(venv_dir, ignore_errors=True)
manifest_path.unlink(missing_ok=True)
uv = Path(args.uv)
if not uv.is_file():
resolved = shutil.which(args.uv)
uv = Path(resolved) if resolved else uv
uv_missing = not uv.is_file()
python = environment_python(runtime_dir)
sync_bundled_vendor(
python = venv_python(venv_dir)
vendor_temp = tempfile.TemporaryDirectory(prefix="snapform_vendor_")
vendor_dir = Path(vendor_temp.name)
copy_vendor_inputs(
bundled_vendor_directory(),
Path(args.runtime_dir).parent / "vendor",
vendor_dir,
)
stages = install_stages(args, runtime_dir, python, uv)
runtime_stages = [stage for stage in stages if stage[3] != "download"]
stages = install_stages(args, venv_dir, python, uv, vendor_dir)
env_stages = [stage for stage in stages if stage[3] != "download"]
model_stages = [stage for stage in stages if stage[3] == "download"]
model_span = (
(MODELS_SLICE_END - MODELS_SLICE_START) / len(model_stages)
@@ -737,13 +748,13 @@ def install(args):
else 0.0
)
installed_packages = []
component_stages = runtime_stages[2:]
component_stages = env_stages[2:]
component_ranges = (
[XFORMERS_SLICE, AI_COMPONENTS_SLICE]
if len(component_stages) == 2
else [(XFORMERS_SLICE[0], AI_COMPONENTS_SLICE[1])]
)
for index, (label, command, environment, kind) in enumerate(runtime_stages):
for index, (label, command, environment, kind) in enumerate(env_stages):
if index == 0:
stage = 1
stage_start, stage_end = ENVIRONMENT_SLICE
@@ -779,21 +790,21 @@ def install(args):
)
if isinstance(stage_installed, list):
installed_packages.extend(stage_installed)
vendor_temp.cleanup()
remove_legacy_vendor(venv_dir.parent)
devices = installed_devices(args)
runtime_dir.mkdir(parents=True, exist_ok=True)
ready_path = runtime_dir / ENVIRONMENT_MANIFEST_FILENAME
venv_dir.mkdir(parents=True, exist_ok=True)
manifest_path.parent.mkdir(parents=True, exist_ok=True)
ready_path = manifest_path
ready_path.write_text(
json.dumps(
{
"environment_manifest_revision": (
ENVIRONMENT_MANIFEST_REVISION
),
"python": "3.12",
"torch": TORCH_VERSION,
"xformers": XFORMERS_VERSION,
"fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION,
"backend": args.backend,
"torch_variant": args.torch_variant,
"devices": devices,
"da3_commit": args.da3_commit,
"ben2_commit": args.ben2_commit,
@@ -834,22 +845,23 @@ def install(args):
write_status(
args.status,
1.0,
f"Backend Environment is ready; retry {labels} from Models",
f"Environment is installed; retry {labels} from Models",
)
return
write_status(
args.status,
1.0,
"Backend Environment and recommended models ready",
"Environment and recommended models ready",
)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--uv", required=True)
parser.add_argument("--runtime-dir", required=True)
parser.add_argument("--venv-dir", required=True)
parser.add_argument("--manifest", required=True)
parser.add_argument("--models-dir", required=True)
parser.add_argument("--backend", choices=tuple(TORCH_INDEXES), required=True)
parser.add_argument("--torch-variant", choices=tuple(TORCH_INDEXES), required=True)
parser.add_argument("--da3-commit", required=True)
parser.add_argument("--ben2-commit", required=True)
parser.add_argument("--status", required=True)
@@ -865,7 +877,7 @@ def main():
write_status(
args.status,
1.0,
"Backend Environment installation failed",
"Environment installation failed",
str(error),
)
raise

View File

@@ -83,7 +83,7 @@ _ENUM_ITEMS_CACHE = {}
def enum_items(supports_gs=None, show_download_status=False):
model_ready = None
if show_download_status:
from .runtime import model_ready
from .client import model_ready
models = [
model

View File

@@ -6,7 +6,7 @@ from pathlib import Path
import bpy
from . import runtime
from . import client
from .models import (
BEN2_MODEL,
get_downloadable_model,
@@ -95,7 +95,7 @@ class ActiveJob:
mesh_create_camera: bool = False
temporary_input_path: Path | None = None
cancelled: bool = False
uses_host: bool = False
uses_service: bool = False
started: bool = False
start_error: str = ""
start_thread: threading.Thread | None = None
@@ -114,9 +114,11 @@ def cancel_active_job():
_active_job.cancelled = True
if not getattr(_active_job, "started", True):
return
if _active_job.uses_host:
if _active_job.uses_service:
try:
runtime.cancel_host_job(_active_job.status_path.stem)
client.cancel_service_job(
client.job_id_from_status(_active_job.status_path)
)
except RuntimeError:
pass
elif _active_job.process is not None and _active_job.process.poll() is None:
@@ -150,7 +152,7 @@ class RunJob(bpy.types.Operator):
job_type: bpy.props.EnumProperty(
items=(
("INSTALL_RUNTIME", "Install Backend", ""),
("INSTALL_ENV", "Install Environment", ""),
("DOWNLOAD_MODEL", "Download Model", ""),
(RECONSTRUCTION_JOB, "Export DA3 Result", ""),
(IMAGE_DEPTH_JOB, "Convert Image to Depth Plane", ""),
@@ -219,10 +221,10 @@ class RunJob(bpy.types.Operator):
settings = context.scene.snapform_settings
try:
command, paths = self._prepare_command(context, settings)
uses_host = isinstance(command, runtime.HostJobRequest)
log_path = runtime.host_log_path() if uses_host else paths["log"]
working_directory = runtime.storage_root()
creation_flags = runtime.creation_flags()
uses_service = isinstance(command, client.JobRequest)
log_path = client.service_log_path() if uses_service else paths["log"]
working_directory = client.storage_root()
creation_flags = client.creation_flags()
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
@@ -262,7 +264,7 @@ class RunJob(bpy.types.Operator):
and self.snapform_delete_input
else None
),
uses_host=uses_host,
uses_service=uses_service,
)
_active_job = job
settings.status = self._starting_message(settings)
@@ -292,11 +294,11 @@ class RunJob(bpy.types.Operator):
job.started = True
return
try:
if getattr(job, "uses_host", False):
runtime.submit_host_job(command)
if getattr(job, "uses_service", False):
client.submit_service_job(command)
else:
if job.kind == "INSTALL_RUNTIME":
runtime.stop_host()
if job.kind == "INSTALL_ENV":
client.stop_service()
job.log_file = job.log_path.open("w", encoding="utf-8")
job.process = subprocess.Popen(
command,
@@ -307,8 +309,10 @@ class RunJob(bpy.types.Operator):
)
job.started = True
if job.cancelled:
if job.uses_host:
runtime.cancel_host_job(job.status_path.stem)
if job.uses_service:
client.cancel_service_job(
client.job_id_from_status(job.status_path)
)
elif job.process is not None:
job.process.terminate()
except (OSError, RuntimeError) as error:
@@ -318,12 +322,12 @@ class RunJob(bpy.types.Operator):
job.start_error = str(error)
def _prepare_command(self, _context, settings):
paths = runtime.new_job_paths(self.job_type)
paths = client.new_job_paths(self.job_type)
if (
getattr(self, "snapform_delete_input", False)
and self.job_type in IMAGE_EMPTY_JOBS
):
archived_input = runtime.move_input_into_job(
archived_input = client.move_input_into_job(
self.snapform_input_path,
paths["result"],
)
@@ -332,18 +336,18 @@ class RunJob(bpy.types.Operator):
remove_background = bool(
getattr(self, "snapform_remove_background", False)
)
if self.job_type == "INSTALL_RUNTIME":
if self.job_type == "INSTALL_ENV":
self._require_online_access()
return runtime.install_runtime_command(paths["status"]), paths
return client.install_env_command(paths["status"]), paths
if not runtime.runtime_ready():
raise RuntimeError("Install the SnapForm Backend first")
if not client.env_installed():
raise RuntimeError("Install the SnapForm Environment first")
if self.job_type == "DOWNLOAD_MODEL":
self._require_online_access()
model = get_downloadable_model(self.model_key)
return (
runtime.download_model_command(model.key, paths["status"]),
client.download_model_command(model.key, paths["status"]),
paths,
)
@@ -372,10 +376,10 @@ class RunJob(bpy.types.Operator):
raise ValueError(
"BEN2 requires an image, video, or directory of image frames"
)
if not runtime.model_ready(BEN2_MODEL.key):
if not client.model_ready(BEN2_MODEL.key):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
return (
runtime.remove_background_request(
client.remove_background_request(
settings,
input_path,
paths["status"],
@@ -388,7 +392,7 @@ class RunJob(bpy.types.Operator):
)
if self.job_type == ORIGINAL_FRAMES_JOB:
return (
runtime.extract_frames_request(
client.extract_frames_request(
settings,
input_path,
paths["status"],
@@ -416,12 +420,12 @@ class RunJob(bpy.types.Operator):
use_ray_pose = self.snapform_use_ray_pose
process_res = self.snapform_process_res
model = get_model(model_key)
if not runtime.model_ready(model_key):
if not client.model_ready(model_key):
raise RuntimeError(f"Download {model.label} first")
if (
self.job_type in DA3_IMAGE_JOBS
and remove_background
and not runtime.model_ready(BEN2_MODEL.key)
and not client.model_ready(BEN2_MODEL.key)
):
raise RuntimeError(f"Download {BEN2_MODEL.label} first")
if export_format == "gs_ply" and not model.supports_gs:
@@ -429,7 +433,7 @@ class RunJob(bpy.types.Operator):
"PLY export requires a DA3 Giant or Nested Giant model"
)
return (
runtime.reconstruct_request(
client.reconstruct_request(
settings,
export_format,
paths["status"],
@@ -456,7 +460,7 @@ class RunJob(bpy.types.Operator):
def _task_device(self, settings):
"""Debug panels use the scene device; production tasks use Preferences."""
if self.job_type in PRODUCTION_AI_JOBS:
return runtime.production_device()
return client.production_device()
return settings.device
def _require_online_access(self):
@@ -467,7 +471,7 @@ class RunJob(bpy.types.Operator):
def _starting_message(self, settings):
labels = {
"INSTALL_RUNTIME": "Starting Backend installation",
"INSTALL_ENV": "Starting Environment installation",
"DOWNLOAD_MODEL": "Starting model download",
}
if self.job_type == RECONSTRUCTION_JOB:
@@ -516,14 +520,14 @@ class RunJob(bpy.types.Operator):
return {"CANCELLED"}
if not job.started:
return {"RUNNING_MODAL"}
if job.uses_host:
status = runtime.read_status(job.status_path) or {}
if job.uses_service:
status = client.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"}:
if client.service_status().get("state") in {"ERROR", "STOPPED"}:
return self._finish(context, job, 1)
return {"RUNNING_MODAL"}
return_code = job.process.poll()
@@ -532,7 +536,7 @@ class RunJob(bpy.types.Operator):
return self._finish(context, job, return_code)
def _update_status(self, context, job):
status = runtime.read_status(job.status_path)
status = client.read_status(job.status_path)
if status is None:
return
settings = context.scene.snapform_settings
@@ -593,28 +597,28 @@ class RunJob(bpy.types.Operator):
if job.cancelled:
settings.status = "Task cancelled"
settings.progress = 0.0
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
if getattr(job, "uses_service", False):
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
status = runtime.read_status(job.status_path) or {}
status = client.read_status(job.status_path) or {}
error = status.get("error")
if return_code != 0 or error:
settings.status = "Task failed"
settings.progress = 0.0
fallback_error = (
"Backend Service task failed"
if job.uses_host
"Service task failed"
if job.uses_service
else "Task subprocess failed"
)
self.report(
{"ERROR"},
f"{error or fallback_error}; log: {job.log_path}",
)
if job.uses_host:
runtime.mark_host_job_complete(job.status_path.stem)
if job.uses_service:
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -644,7 +648,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -659,7 +663,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Surface import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -674,7 +678,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Depth import failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -689,7 +693,7 @@ class RunJob(bpy.types.Operator):
settings.status = "Image update failed"
settings.progress = 0.0
self.report({"ERROR"}, str(import_error))
runtime.mark_host_job_complete(job.status_path.stem)
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"CANCELLED"}
@@ -701,8 +705,8 @@ class RunJob(bpy.types.Operator):
job.status_path.unlink(missing_ok=True)
except OSError:
pass
if getattr(job, "uses_host", False):
runtime.mark_host_job_complete(job.status_path.stem)
if getattr(job, "uses_service", False):
client.mark_service_job_complete(client.job_id_from_status(job.status_path))
close_active_job()
self._redraw_ui(context, force=True)
return {"FINISHED"}
@@ -836,17 +840,17 @@ class CancelJob(bpy.types.Operator):
return {"FINISHED"}
class ControlHost(bpy.types.Operator):
bl_idname = "snapform.control_host"
bl_label = "Control SnapForm Backend Service"
class ControlService(bpy.types.Operator):
bl_idname = "snapform.control_service"
bl_label = "Control SnapForm Service"
action: bpy.props.EnumProperty(
items=(
("START", "Start", "Start or connect to the Backend Service"),
("START", "Start", "Start or connect to the Service"),
("CLEAR", "Clear Models", "Release cached AI models"),
("RESTART", "Restart", "Restart the Backend Service"),
("STOP", "Stop", "Stop the Backend Service"),
("LOG", "Open Log", "Open the Backend Service log"),
("RESTART", "Restart", "Restart the Service"),
("STOP", "Stop", "Stop the Service"),
("LOG", "Open Log", "Open the Service log"),
),
options={"HIDDEN", "SKIP_SAVE"},
)
@@ -854,23 +858,23 @@ class ControlHost(bpy.types.Operator):
def execute(self, _context):
try:
if self.action == "START":
runtime.enable_host()
message = "Backend Service is ready"
client.enable_service()
message = "Service is ready"
elif self.action == "CLEAR":
runtime.clear_host_models()
client.clear_service_models()
message = "AI model cache is cleared"
elif self.action == "RESTART":
runtime.restart_host()
message = "Backend Service is restarting"
client.restart_service()
message = "Service is restarting"
elif self.action == "STOP":
runtime.stop_host()
message = "Backend Service is stopping"
client.stop_service()
message = "Service is stopping"
else:
log_path = runtime.host_log_path()
log_path = client.service_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 Backend Service log")
raise RuntimeError("Unable to open the Service log")
return {"FINISHED"}
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))

View File

@@ -2,13 +2,13 @@ import bpy
from .models import BEN2_MODEL, MODELS, selected_model_key
from .operators import active_job
from . import runtime
from .runtime import model_ready, runtime_ready
from . import client
from .client import model_ready, env_installed
class HostPanel(bpy.types.Panel):
bl_label = "AI Backend"
bl_idname = "SNAPFORM_PT_host"
class AiServicePanel(bpy.types.Panel):
bl_label = "AI Service"
bl_idname = "SNAPFORM_PT_service"
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
@@ -16,23 +16,23 @@ class HostPanel(bpy.types.Panel):
def draw(self, _context):
layout = self.layout
running = active_job() is not None or runtime.host_busy()
if not runtime_ready():
running = active_job() is not None or client.service_busy()
if not env_installed():
layout.label(
text="Backend Environment is not installed",
text="Environment is not installed",
icon="ERROR",
)
install_row = layout.row()
install_row.enabled = not running
operator = install_row.operator(
"snapform.run_job",
text="Install Backend",
text="Install Environment",
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
operator.job_type = "INSTALL_ENV"
return
status = runtime.host_status()
status = client.service_status()
state = status.get("state", "STOPPED")
icons = {
"READY": "RADIOBUT_ON",
@@ -56,20 +56,20 @@ class HostPanel(bpy.types.Panel):
controls = row.row(align=True)
controls.enabled = status.get("clients", 0) <= 1
operator = controls.operator(
"snapform.control_host",
"snapform.control_service",
text="",
icon="FILE_REFRESH",
)
operator.action = "RESTART"
operator = controls.operator(
"snapform.control_host",
"snapform.control_service",
text="",
icon="PAUSE",
)
operator.action = "STOP"
elif state == "STOPPED":
operator = row.operator(
"snapform.control_host",
"snapform.control_service",
text="Start",
icon="PLAY",
)
@@ -81,7 +81,7 @@ class HostPanel(bpy.types.Panel):
)
if state == "ERROR":
operator = row.operator(
"snapform.control_host",
"snapform.control_service",
text="Start",
icon="PLAY",
)
@@ -101,7 +101,7 @@ class HostPanel(bpy.types.Panel):
clear = cached_row.row(align=True)
clear.enabled = state == "READY"
operator = clear.operator(
"snapform.control_host",
"snapform.control_service",
text="",
icon="TRASH",
)
@@ -111,20 +111,19 @@ class HostPanel(bpy.types.Panel):
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
bl_order = 1
@classmethod
def poll(cls, _context):
return runtime_ready()
return env_installed()
def draw(self, _context):
layout = self.layout
running = active_job() is not None or runtime.host_busy()
running = active_job() is not None or client.service_busy()
for model in MODELS.values():
self._draw_model(layout, model, running)
@@ -137,7 +136,7 @@ class ModelsPanel(bpy.types.Panel):
row.label(text="Ready", icon="CHECKMARK")
else:
download_row = row.row(align=True)
download_row.enabled = runtime_ready() and not running
download_row.enabled = env_installed() and not running
operator = download_row.operator(
"snapform.run_job",
text="Download",
@@ -154,7 +153,7 @@ class DebugPanel(bpy.types.Panel):
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 1
bl_order = 2
def draw(self, context):
layout = self.layout
@@ -226,7 +225,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
export_row = layout.row()
export_row.enabled = (
active_job() is None
and not runtime.host_busy()
and not client.service_busy()
and model_ready(selected_model_key(settings))
)
operator = export_row.operator(
@@ -245,7 +244,7 @@ class Ben2RemovalDebugPanel(bpy.types.Panel):
bl_region_type = "UI"
bl_category = "SnapForm"
bl_options = {"DEFAULT_CLOSED"}
bl_order = 1
bl_order = 2
def draw(self, context):
layout = self.layout
@@ -255,7 +254,7 @@ class Ben2RemovalDebugPanel(bpy.types.Panel):
background_row = layout.row()
background_row.enabled = (
active_job() is None
and not runtime.host_busy()
and not client.service_busy()
and model_ready(BEN2_MODEL.key)
)
operator = background_row.operator(

View File

@@ -2,7 +2,7 @@ from pathlib import Path
import bpy
from .runtime import device_enum_items
from .client import device_enum_items
class SnapFormPreferences(bpy.types.AddonPreferences):
@@ -11,7 +11,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
storage_root: bpy.props.StringProperty(
name="Storage Root",
description=(
"Shared user-writable storage for Backend Environment, models and jobs; "
"Shared user-writable storage for the environment, models and jobs; "
"all Blender versions can use the same directory"
),
subtype="DIR_PATH",
@@ -29,27 +29,27 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
def draw(self, _context):
from .operators import active_job
from .runtime import runtime_directory, runtime_ready
from .client import venv_directory, env_installed
layout = self.layout
layout.prop(self, "storage_root")
layout.prop(self, "device")
layout.separator()
layout.label(text=f"Backend Environment: {runtime_directory()}")
if runtime_ready():
layout.label(text="Backend Environment is ready", icon="CHECKMARK")
layout.label(text=f"Environment: {venv_directory()}")
if env_installed():
layout.label(text="Environment is installed", icon="CHECKMARK")
else:
install_row = layout.row()
install_row.enabled = active_job() is None
operator = install_row.operator(
"snapform.run_job",
text="Install Backend",
text="Install Environment",
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
operator.job_type = "INSTALL_ENV"
layout.separator()
operator = layout.operator(
"snapform.control_host",
"snapform.control_service",
text="Open Service Log",
icon="TEXT",
)

View File

@@ -1,7 +1,7 @@
import bpy
from .models import gs_model_enum_items, model_enum_items
from .runtime import device_enum_items
from .client import device_enum_items
class SnapFormSettings(bpy.types.PropertyGroup):

View File

@@ -19,12 +19,11 @@ if str(SCRIPT_DIRECTORY) not in sys.path:
import worker
JOB_PROTOCOL_VERSION = 2
CLIENT_TIMEOUT_SECONDS = 10.0
IDLE_SHUTDOWN_SECONDS = 20.0
class HostLock:
class ServiceLock:
def __init__(self, path):
self.path = Path(path)
self.file = None
@@ -159,7 +158,7 @@ class ModelCache:
self.torch.cuda.empty_cache()
class HostState:
class ServiceState:
def __init__(self, token, instance_id):
self.token = token
self.instance_id = instance_id
@@ -212,7 +211,6 @@ class HostState:
with self.lock:
return {
"ready": True,
"job_protocol_version": JOB_PROTOCOL_VERSION,
"instance_id": self.instance_id,
"busy": self.active_job is not None,
"active_job": dict(self.active_job) if self.active_job else None,
@@ -265,7 +263,7 @@ class HostState:
worker.write_status(
status_path,
0.0,
"Backend Service accepted the task",
"Service accepted the task",
state="running",
)
try:
@@ -318,7 +316,7 @@ class HostState:
def start_monitor(self):
self.monitor_thread = threading.Thread(
target=self._monitor,
name="SnapFormHostMonitor",
name="SnapFormServiceMonitor",
daemon=True,
)
self.monitor_thread.start()
@@ -363,7 +361,7 @@ class HostState:
def create_app(token, instance_id="test"):
from fastapi import BackgroundTasks, FastAPI, Header, HTTPException
state = HostState(token, instance_id)
state = ServiceState(token, instance_id)
@asynccontextmanager
async def lifespan(_app):
@@ -372,17 +370,17 @@ def create_app(token, instance_id="test"):
state.close()
app = FastAPI(
title="SnapForm Backend Service",
title="SnapForm Service",
docs_url=None,
redoc_url=None,
openapi_url=None,
lifespan=lifespan,
)
app.state.snapform_host_state = state
app.state.snapform_service_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")
raise HTTPException(status_code=401, detail="Invalid service token")
@app.get("/health")
def health(x_snapform_token: str = Header(default="")):
@@ -445,7 +443,7 @@ def create_app(token, instance_id="test"):
command,
parameters["status"],
):
raise HTTPException(status_code=409, detail="Backend Service is busy")
raise HTTPException(status_code=409, detail="Service is busy")
worker.write_status(
parameters["status"],
0.0,
@@ -485,11 +483,11 @@ def create_app(token, instance_id="test"):
def clear_models(x_snapform_token: str = Header(default="")):
authorize(x_snapform_token)
if not state.clear_models():
raise HTTPException(status_code=409, detail="Backend Service is busy")
raise HTTPException(status_code=409, detail="Service is busy")
return {"state": "cleared"}
@app.post("/shutdown", status_code=202)
def shutdown_host(
def shutdown_service(
payload: dict,
x_snapform_token: str = Header(default=""),
):
@@ -498,7 +496,7 @@ def create_app(token, instance_id="test"):
if result == "shared":
raise HTTPException(status_code=409, detail="Other Blender clients are connected")
if result == "busy":
raise HTTPException(status_code=409, detail="Backend Service is busy")
raise HTTPException(status_code=409, detail="Service is busy")
if result != "accepted":
raise HTTPException(status_code=404, detail="Client is not registered")
return {"state": "stopping"}
@@ -524,7 +522,6 @@ def write_discovery(args):
temporary.write_text(
json.dumps(
{
"job_protocol_version": JOB_PROTOCOL_VERSION,
"port": args.port,
"token": args.token,
"pid": os.getpid(),
@@ -551,8 +548,8 @@ def main():
import uvicorn
args = parse_args()
host_lock = HostLock(args.lock)
if not host_lock.acquire():
service_lock = ServiceLock(args.lock)
if not service_lock.acquire():
return
try:
app = create_app(args.token, args.instance_id)
@@ -566,7 +563,7 @@ def main():
server = uvicorn.Server(config)
def watch_shutdown():
app.state.snapform_host_state.shutdown_event.wait()
app.state.snapform_service_state.shutdown_event.wait()
server.should_exit = True
threading.Thread(target=watch_shutdown, daemon=True).start()
@@ -574,7 +571,7 @@ def main():
server.run()
finally:
remove_discovery(args.discovery, args.instance_id)
host_lock.release()
service_lock.release()
if __name__ == "__main__":

View File

@@ -1282,7 +1282,7 @@ def infer_foreground_frames(
def remove_background(args, model_cache=None, cancel_check=None):
total_started = time.perf_counter()
write_status(args.status, 0.02, "Starting Backend Environment")
write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check)
if model_cache is None:
@@ -1501,7 +1501,7 @@ def flatten_exported_paths(
def reconstruct(args, model_cache=None, cancel_check=None):
write_status(args.status, 0.02, "Starting Backend Environment")
write_status(args.status, 0.02, "Starting Environment")
check_cancelled(cancel_check)
if model_cache is None: