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

Add ONNX model inference implementations for MoGe-2, BEN2, and DA3

Browse Source 
- Introduced `onnx_moge2.py` for MoGe-2 ONNX inference with NumPy post-processing.
- Added `onnx_ben2.py` for BEN2 ONNX inference with Pillow and NumPy preprocessing.
- Created `onnx_da3.py` for Depth Anything 3 inference with ONNX Runtime and NumPy.
- Implemented shared ONNX Runtime session helpers in `onnx_runtime.py`.
- Developed unit tests for server functionalities in `test_server.py`, including model caching and job management.
- Enhanced job server to handle frame extraction without requiring models.
This commit is contained in:
马南
2026-08-17 10:37:55 +08:00
parent b3a07cde32
commit c5b77c5c85
58 changed files with 5840 additions and 9931 deletions
Download Patch File Download Diff File Expand all files Collapse all files

39
src/snapform/__init__.py
View File

@@ -1,29 +1,21 @@
import bpy
from . import client
from .blender import depth_plane as image_depth
from .blender import surface as image_mesh
from .blender.depth_plane import (
ApplyRemovedBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
BuildMoge2DepthPlane,
ConvertImageToDepthPlane,
ConvertToDepthPlane,
RemoveImageBackground,
SnapFormImageMenu,
draw_image_context_menu,
)
from .blender.surface import BuildImageMesh, ConvertImageToMesh
from .blender.surface import ConvertToSurface
from .operators import (
CancelJob,
ControlService,
JOB_OPERATOR_CLASSES,
cancel_active_job,
close_active_job,
)
from .job_runtime import runtime
from .panel import (
Ben2RemovalDebugPanel,
AiServicePanel,
ServerPanel,
Da3OutputDebugPanel,
DebugPanel,
ModelsPanel,
@@ -31,25 +23,17 @@ from .panel import (
)
from .preferences import SnapFormPreferences
from .properties import SnapFormSettings
from .status_bar import draw_status_bar
CLASSES = (
SnapFormPreferences,
SnapFormSettings,
*JOB_OPERATOR_CLASSES,
CancelJob,
ControlService,
SnapFormImageMenu,
ApplyRemovedBackground,
RemoveImageBackground,
BuildDepthPlane,
BuildOriginalDepthPlane,
BuildMoge2DepthPlane,
ConvertImageToDepthPlane,
BuildImageMesh,
ConvertImageToMesh,
AiServicePanel,
ConvertToDepthPlane,
ConvertToSurface,
ServerPanel,
ModelsPanel,
DebugPanel,
Moge2ApiDebugPanel,
@@ -64,20 +48,15 @@ def register():
bpy.types.Scene.snapform_settings = bpy.props.PointerProperty(
type=SnapFormSettings,
)
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)
client.schedule_service_start()
runtime.register()
def unregister():
cancel_active_job()
close_active_job()
client.unschedule_service()
client.detach_service_client()
runtime.unregister()
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)
if hasattr(bpy.types.Scene, "snapform_settings"):
del bpy.types.Scene.snapform_settings
for class_type in reversed(CLASSES):

1
src/snapform/ai_service/__init__.py
View File

@@ -1 +0,0 @@
"""SnapForm AI Service implementation without Blender dependencies."""

2277
src/snapform/ai_service/jobs.py
View File

File diff suppressed because it is too large Load Diff

145
src/snapform/ai_service/server.py
View File

@@ -1,145 +0,0 @@
import argparse
import sys
import threading
from pathlib import Path
from types import SimpleNamespace
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
ADDON_DIRECTORY = SCRIPT_DIRECTORY.parent
for import_directory in (SCRIPT_DIRECTORY, ADDON_DIRECTORY):
if str(import_directory) not in sys.path:
sys.path.insert(0, str(import_directory))
import jobs as worker
from model_manager import ModelManager, ModelSessionCache
from blendjob.server import JobCancelled, JobServer, watch_parent
def create_server(models_directory=None, jobs_directory=None):
"""Create the SnapForm-specific server and register its job types."""
server = JobServer("SnapForm Job Server", jobs_directory=jobs_directory)
models = server.add_resource(
"models",
ModelManager(models_directory or Path.cwd() / "models"),
)
def run_worker(context, parameters, function, *, uses_models=True):
parameters = dict(parameters)
parameters["output_dir"] = str(context.directory)
models = context.resource("models") if uses_models else None
model_key = parameters.pop("model", "")
if model_key:
parameters["model_dir"] = str(models.directory(model_key))
background_model = parameters.pop("remove_background_model", "")
if background_model:
parameters["remove_background_model_dir"] = str(
models.directory(background_model)
)
args = SimpleNamespace(**parameters)
args.status = context.update
options = {"cancel_check": context.is_cancelled}
if uses_models:
options["model_cache"] = models
try:
function(args, **options)
except worker.TaskCancelled as error:
raise JobCancelled(str(error))
@server.job("remove-background")
def remove_background(context, parameters):
run_worker(context, parameters, worker.remove_background)
@server.job("debug-ben2")
def debug_ben2(context, parameters):
run_worker(context, parameters, worker.remove_background)
@server.job("extract-frames")
def extract_frames(context, parameters):
run_worker(
context,
parameters,
worker.extract_frames,
uses_models=False,
)
@server.job("generate-da3-depth")
def generate_da3_depth(context, parameters):
run_worker(context, parameters, worker.reconstruct)
@server.job("generate-da3-surface")
def generate_da3_surface(context, parameters):
run_worker(context, parameters, worker.reconstruct)
@server.job("debug-da3")
def debug_da3(context, parameters):
run_worker(context, parameters, worker.reconstruct)
@server.job("generate-moge2-depth")
def generate_moge2_depth(context, parameters):
run_worker(context, parameters, worker.run_moge)
@server.job("generate-moge2-surface")
def generate_moge2_surface(context, parameters):
parameters = {**parameters, "surface_mesh": True}
run_worker(context, parameters, worker.run_moge)
@server.job("debug-moge2")
def debug_moge2(context, parameters):
parameters = dict(parameters)
parameters["debug_export_format"] = parameters.pop("export_format")
run_worker(context, parameters, worker.run_moge)
@server.job("download-model")
def download_model(context, parameters):
context.resource("models").download(context, parameters["model"])
return server
def create_app(instance_id="test", server=None):
"""Create the dedicated SnapForm FastAPI app."""
server = server or create_server()
return server.create_app(instance_id)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--instance-id", required=True)
parser.add_argument("--parent-pid", type=int, required=True)
return parser.parse_args()
def main():
import uvicorn
args = parse_args()
job_server = create_server()
app = create_app(args.instance_id, job_server)
config = uvicorn.Config(
app,
host=args.host,
port=args.port,
access_log=False,
log_level="warning",
)
http_server = uvicorn.Server(config)
def stop_http_server():
job_server.shutdown_event.wait()
http_server.should_exit = True
threading.Thread(target=stop_http_server, daemon=True).start()
threading.Thread(
target=watch_parent,
args=(args.parent_pid, job_server.shutdown_event),
name="SnapFormParentMonitor",
daemon=True,
).start()
http_server.run()
if __name__ == "__main__":
main()

388
src/snapform/blender/depth_plane.py
View File

@@ -1,5 +1,6 @@
import re
import shutil
import tempfile
from pathlib import Path
import bpy
@@ -53,7 +54,7 @@ VIDEO_NAME_SUFFIXES = (
def image_depth_model_items(owner, context):
from ..models import depth_plane_model_enum_items
from ..properties import depth_plane_model_enum_items
return depth_plane_model_enum_items(owner, context)
@@ -218,19 +219,17 @@ def process_resolution_scale_property():
def draw_process_resolution(layout, owner, image):
process_res = scaled_process_resolution(
image,
owner.snapform_process_res_factor,
owner.process_res_factor,
)
layout.prop(
owner,
"snapform_process_res_factor",
"process_res_factor",
text=f"Resolution Scale ({process_res} px)",
)
def save_packed_image(image, scene):
from ..client import temporary_input_path
output_path = temporary_input_path()
output_path = Path(tempfile.mkdtemp(prefix="snapform-input-")) / "input.png"
image_settings = scene.render.image_settings
original_file_format = image_settings.file_format
try:
@@ -286,9 +285,7 @@ def prepare_image_input(image, scene):
return source_path, False
sequence_frames = collect_sequence_frames(image)
if sequence_frames:
from ..client import temporary_input_directory
temporary_directory = Path(temporary_input_directory())
temporary_directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
try:
for index, frame_path in enumerate(sequence_frames):
suffix = frame_path.suffix.lower() or ".jpg"
@@ -1868,90 +1865,6 @@ def build_depth_plane(
)
class BuildDepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_depth_plane"
bl_label = "Build Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_data_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
npz_path = Path(self.snapform_npz_path)
if not npz_path.is_file():
self.report(
{"ERROR"},
f"Depth result does not exist: {npz_path}",
)
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_depth_plane(
context,
source_object,
npz_path,
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
depth_name=self.snapform_depth_data_name,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
def build_moge2_depth_plane(
context,
source_object,
@@ -2048,177 +1961,6 @@ def build_moge2_depth_plane(
)
class BuildMoge2DepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_moge2_depth_plane"
bl_label = "Build MoGe-2 Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(options={"HIDDEN"})
snapform_npz_path: bpy.props.StringProperty(options={"HIDDEN"})
snapform_background_path: bpy.props.StringProperty(options={"HIDDEN"})
snapform_color_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
snapform_data_frames_dir: bpy.props.StringProperty(options={"HIDDEN"})
snapform_delete_source: bpy.props.BoolProperty(default=True, options={"HIDDEN"})
snapform_unshaded: bpy.props.BoolProperty(default=False, options={"HIDDEN"})
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report({"ERROR"}, "The source image empty is no longer available")
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
build_moge2_depth_plane(
context,
source_object,
Path(self.snapform_npz_path),
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
color_image=color_image,
foreground_image=foreground_image,
data_frames_directory=(
self.snapform_data_frames_dir or None
),
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class BuildOriginalDepthPlane(bpy.types.Operator):
bl_idname = "snapform.build_original_depth_plane"
bl_label = "Build Original Depth Plane"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_background_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_color_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_frames_dir: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_depth_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
foreground_image = None
if self.snapform_background_path:
foreground_image = create_removed_background_image(
source_object,
self.snapform_background_path,
assign_to_source=False,
)
color_image = None
# The foreground result already carries the color, so the original
# frame sequence is unused when background removal is available.
if self.snapform_color_frames_dir and foreground_image is None:
color_image = create_sequence_image(
f"{image_base_name(source_object.data)}.frame",
self.snapform_color_frames_dir,
)
depth_image = None
if self.snapform_depth_npz_path:
depth_arrays = extract_depth_arrays(
self.snapform_depth_npz_path
)
if len(depth_arrays) == 1:
depth_image = create_depth_image(
depth_data_name(source_object),
depth_arrays[0],
)
else:
depth_image = create_depth_sequence_image(
depth_data_name(source_object),
depth_arrays,
context.scene,
self.snapform_depth_frames_dir,
)
build_image_depth_plane(
context,
source_object,
self.snapform_delete_source,
unshaded=self.snapform_unshaded,
color_image=color_image,
foreground_image=foreground_image,
depth_frames_directory=(
self.snapform_depth_frames_dir or None
),
depth_image=depth_image,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class ApplyRemovedBackground(bpy.types.Operator):
bl_idname = "snapform.apply_removed_background"
bl_label = "Apply Removed Background"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_output_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, _context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report(
{"ERROR"},
"The source image empty is no longer available",
)
return {"CANCELLED"}
try:
create_removed_background_image(
source_object,
self.snapform_output_path,
)
except RuntimeError as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class RemoveImageBackground(bpy.types.Operator):
bl_idname = "snapform.remove_image_background"
bl_label = "Remove Background"
@@ -2227,15 +1969,13 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from ..operators import active_job
from ..client import service_busy, model_ready
from ..models import BEN2_MODEL
from ..job_runtime import runtime
from ..properties import BEN2_MODEL_KEY, model_ready
return (
is_image_empty(context.object)
and active_job() is None
and not service_busy()
and model_ready(BEN2_MODEL.key)
and not runtime.server_busy()
and model_ready(BEN2_MODEL_KEY)
)
def execute(self, context):
@@ -2255,9 +1995,9 @@ class RemoveImageBackground(bpy.types.Operator):
return {"CANCELLED"}
result = bpy.ops.snapform.process_image_background(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_input=delete_input,
input_path=str(source_path),
source_object_name=source_object.name,
delete_input=delete_input,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2266,42 +2006,42 @@ class RemoveImageBackground(bpy.types.Operator):
return {"CANCELLED"}
class ConvertImageToDepthPlane(bpy.types.Operator):
bl_idname = "snapform.convert_image_to_depth_plane"
class ConvertToDepthPlane(bpy.types.Operator):
bl_idname = "snapform.convert_to_depth_plane"
bl_label = "Convert to Depth Plane"
bl_description = "Convert an image to a displaced depth plane"
bl_options = {"REGISTER"}
snapform_unshaded: bpy.props.BoolProperty(
unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the background during conversion with BEN2",
default=False,
)
snapform_model: bpy.props.EnumProperty(
model: bpy.props.EnumProperty(
name="Depth From",
description="Model or original image used to build depth",
items=image_depth_model_items,
default=0,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty(
process_res_factor: process_resolution_scale_property()
resolution_level: bpy.props.IntProperty(
name="Resolution Level",
description="MoGe-2 inference resolution level; higher is finer and slower",
default=5,
min=0,
max=9,
)
snapform_keep_original: bpy.props.BoolProperty(
keep_original: bpy.props.BoolProperty(
name="Keep Original",
description="Keep the source Image Empty after conversion",
default=False,
@@ -2309,13 +2049,11 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
@classmethod
def poll(cls, context):
from ..operators import active_job
from ..client import service_busy
from ..job_runtime import runtime
return (
is_image_empty(context.object)
and active_job() is None
and not service_busy()
and not runtime.server_busy()
)
def invoke(self, context, _event):
@@ -2323,15 +2061,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
def draw(self, context):
layout = self.layout
layout.prop(self, "snapform_model")
if self.snapform_model in {"SMALL", "BASE"}:
layout.prop(self, "model")
if self.model in {"SMALL", "BASE"}:
draw_process_resolution(layout, self, context.object.data)
elif self.snapform_model.startswith("MOGE2_"):
layout.prop(self, "snapform_moge_resolution_level")
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded")
elif self.model.startswith("MOGE2_"):
layout.prop(self, "resolution_level")
layout.prop(self, "remove_background")
layout.prop(self, "unshaded")
layout.separator(type="LINE")
layout.prop(self, "snapform_keep_original")
layout.prop(self, "keep_original")
def execute(self, context):
source_object = context.object
@@ -2339,7 +2077,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
context.scene.snapform_settings.video_frames = (
default_extract_frames(source_object)
)
if self.snapform_model == "ORIGINAL":
if self.model == "ORIGINAL":
try:
source_path, delete_input = prepare_image_input(
source_object.data,
@@ -2350,15 +2088,15 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
run_original = (
bpy.ops.snapform.generate_original_background
if self.snapform_remove_background
if self.remove_background
else bpy.ops.snapform.extract_original_frames
)
result = run_original(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
input_path=str(source_path),
source_object_name=source_object.name,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2366,7 +2104,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
remove_temporary_input(source_path)
return {"CANCELLED"}
if self.snapform_model.startswith("MOGE2_"):
if self.model.startswith("MOGE2_"):
try:
source_path, delete_input = prepare_image_input(
source_object.data,
@@ -2376,16 +2114,14 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
result = bpy.ops.snapform.generate_moge2_depth(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_moge_resolution_level=(
self.snapform_moge_resolution_level
),
snapform_model_key=self.snapform_model,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=self.snapform_remove_background,
input_path=str(source_path),
source_object_name=source_object.name,
resolution_level=self.resolution_level,
model=self.model,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
remove_background=self.remove_background,
)
if "CANCELLED" not in result:
return {"FINISHED"}
@@ -2397,7 +2133,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
try:
process_res = scaled_process_resolution(
image,
self.snapform_process_res_factor,
self.process_res_factor,
)
except ValueError as error:
self.report({"ERROR"}, str(error))
@@ -2412,16 +2148,16 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
return {"CANCELLED"}
result = bpy.ops.snapform.generate_da3_depth(
snapform_input_path=str(source_path),
snapform_source_object_name=source_object.name,
snapform_depth_data_name=depth_data_name(source_object),
snapform_model_key=self.snapform_model,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
snapform_delete_source=not self.snapform_keep_original,
snapform_delete_input=delete_input,
snapform_unshaded=self.snapform_unshaded,
snapform_remove_background=self.snapform_remove_background,
input_path=str(source_path),
source_object_name=source_object.name,
depth_data_name=depth_data_name(source_object),
model=self.model,
use_ray_pose=self.use_ray_pose,
process_res=process_res,
delete_source=not self.keep_original,
delete_input=delete_input,
unshaded=self.unshaded,
remove_background=self.remove_background,
)
if "CANCELLED" in result:
if delete_input:
@@ -2435,11 +2171,11 @@ class SnapFormImageMenu(bpy.types.Menu):
bl_label = "SnapForm"
def draw(self, _context):
from ..client import env_installed
from ..job_runtime import runtime
layout = self.layout
layout.operator_context = "INVOKE_DEFAULT"
if not env_installed():
if not runtime.environment_ready():
layout.operator(
"snapform.install_environment",
text="Install Environment",
@@ -2447,12 +2183,12 @@ class SnapFormImageMenu(bpy.types.Menu):
)
return
layout.operator(
"snapform.convert_image_to_depth_plane",
"snapform.convert_to_depth_plane",
text="Convert to Depth Plane",
icon="MOD_DISPLACE",
)
layout.operator(
"snapform.convert_image_to_mesh",
"snapform.convert_to_surface",
text="Convert to Surface",
icon="MESH_GRID",
)

129
src/snapform/blender/surface.py
View File

@@ -37,7 +37,7 @@ NODE_ASSET_NAME = "O_SnapForm.blend"
def image_surface_model_items(owner, context):
from ..models import surface_model_enum_items
from ..properties import surface_model_enum_items
return surface_model_enum_items(owner, context)
@@ -554,96 +554,47 @@ def build_image_mesh(
return mesh_object
class BuildImageMesh(bpy.types.Operator):
bl_idname = "snapform.build_image_mesh"
bl_label = "Build Image Surface"
bl_options = {"INTERNAL"}
snapform_source_object_name: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_npz_path: bpy.props.StringProperty(
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_delete_source: bpy.props.BoolProperty(
default=True,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_create_camera: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_unshaded: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
snapform_moge2: bpy.props.BoolProperty(
default=False,
options={"HIDDEN", "SKIP_SAVE"},
)
def execute(self, context):
source_object = bpy.data.objects.get(self.snapform_source_object_name)
if not is_image_empty(source_object):
self.report({"ERROR"}, "The source Image Empty no longer exists")
return {"CANCELLED"}
try:
build_image_mesh(
context,
source_object,
self.snapform_npz_path,
self.snapform_delete_source,
self.snapform_create_camera,
self.snapform_unshaded,
self.snapform_moge2,
)
except (OSError, RuntimeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class ConvertImageToMesh(bpy.types.Operator):
bl_idname = "snapform.convert_image_to_mesh"
class ConvertToSurface(bpy.types.Operator):
bl_idname = "snapform.convert_to_surface"
bl_label = "Convert to Surface"
bl_description = "Convert an Image Empty into a perspective textured surface"
bl_options = {"REGISTER"}
snapform_model: bpy.props.EnumProperty(
model: bpy.props.EnumProperty(
name="Depth From",
description="Model used to estimate image depth",
items=image_surface_model_items,
default=0,
)
snapform_use_ray_pose: bpy.props.BoolProperty(
use_ray_pose: bpy.props.BoolProperty(
name="Use Ray Pose",
description="Legacy task compatibility; ONNX uses the camera decoder",
default=False,
)
snapform_process_res_factor: process_resolution_scale_property()
snapform_moge_resolution_level: bpy.props.IntProperty(
process_res_factor: process_resolution_scale_property()
resolution_level: bpy.props.IntProperty(
name="Resolution Level",
description="MoGe-2 inference resolution level; higher is finer and slower",
default=5,
min=0,
max=9,
)
snapform_create_camera: bpy.props.BoolProperty(
create_camera: bpy.props.BoolProperty(
name="Create Matching Camera",
description="Create a perspective camera matching the source projection",
default=False,
)
snapform_unshaded: bpy.props.BoolProperty(
unshaded: bpy.props.BoolProperty(
name="Unshaded",
description="Display the source color through an Emission shader",
default=False,
)
snapform_remove_background: bpy.props.BoolProperty(
remove_background: bpy.props.BoolProperty(
name="Remove Background",
description="Remove the background during conversion with BEN2",
default=False,
)
snapform_keep_original: bpy.props.BoolProperty(
keep_original: bpy.props.BoolProperty(
name="Keep Original",
description="Keep the source Image Empty after conversion",
default=False,
@@ -651,18 +602,18 @@ class ConvertImageToMesh(bpy.types.Operator):
@classmethod
def poll(cls, context):
from ..models import MODELS, MOGE2_MODELS
from ..operators import active_job
from ..client import service_busy, model_ready
from ..job_runtime import runtime
from ..properties import model_catalog
return (
is_image_empty(context.object)
and video_total_frames(context.object.data) == 1
and active_job() is None
and not service_busy()
and not runtime.server_busy()
and any(
model_ready(model_key)
for model_key in (*MOGE2_MODELS, *MODELS)
model.get("ready")
and model.get("enabled")
and model.get("family") in {"moge2", "da3"}
for model in model_catalog()
)
)
@@ -671,16 +622,16 @@ class ConvertImageToMesh(bpy.types.Operator):
def draw(self, context):
layout = self.layout
layout.prop(self, "snapform_model")
if self.snapform_model.startswith("MOGE2_"):
layout.prop(self, "snapform_moge_resolution_level")
layout.prop(self, "model")
if self.model.startswith("MOGE2_"):
layout.prop(self, "resolution_level")
else:
draw_process_resolution(layout, self, context.object.data)
layout.prop(self, "snapform_create_camera")
layout.prop(self, "snapform_remove_background")
layout.prop(self, "snapform_unshaded")
layout.prop(self, "create_camera")
layout.prop(self, "remove_background")
layout.prop(self, "unshaded")
layout.separator(type="LINE")
layout.prop(self, "snapform_keep_original")
layout.prop(self, "keep_original")
def execute(self, context):
source_object = context.object
@@ -689,10 +640,10 @@ class ConvertImageToMesh(bpy.types.Operator):
return {"CANCELLED"}
try:
process_res = 1
if not self.snapform_model.startswith("MOGE2_"):
if not self.model.startswith("MOGE2_"):
process_res = scaled_process_resolution(
source_object.data,
self.snapform_process_res_factor,
self.process_res_factor,
)
source_path, delete_input = prepare_image_input(
source_object.data,
@@ -702,27 +653,25 @@ class ConvertImageToMesh(bpy.types.Operator):
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
common = {
"snapform_input_path": str(source_path),
"snapform_source_object_name": source_object.name,
"snapform_model_key": self.snapform_model,
"snapform_delete_source": not self.snapform_keep_original,
"snapform_delete_input": delete_input,
"snapform_mesh_create_camera": self.snapform_create_camera,
"snapform_unshaded": self.snapform_unshaded,
"snapform_remove_background": self.snapform_remove_background,
"input_path": str(source_path),
"source_object_name": source_object.name,
"model": self.model,
"delete_source": not self.keep_original,
"delete_input": delete_input,
"create_camera": self.create_camera,
"unshaded": self.unshaded,
"remove_background": self.remove_background,
}
if self.snapform_model.startswith("MOGE2_"):
if self.model.startswith("MOGE2_"):
result = bpy.ops.snapform.generate_moge2_surface(
**common,
snapform_moge_resolution_level=(
self.snapform_moge_resolution_level
),
resolution_level=self.resolution_level,
)
else:
result = bpy.ops.snapform.generate_da3_surface(
**common,
snapform_use_ray_pose=self.snapform_use_ray_pose,
snapform_process_res=process_res,
use_ray_pose=self.use_ray_pose,
process_res=process_res,
)
if "CANCELLED" in result:
if delete_input:

2
src/snapform/blender_manifest.toml
View File

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

38
src/snapform/blendjob/__init__.py
View File

@@ -1,37 +1,21 @@
"""Reusable local HTTP job service for Blender add-ons."""
"""Reusable local HTTP Job Server for Blender add-ons."""
from .client import JobClient, JobRequest
from .client import JobClient, JobResult
from .controller import ServerConnection, ServerController
from .operator import JobOperatorBase, JobOperatorState
from .runtime import BlenderJobRuntime
from .server import JobCancelled, JobContext, JobServer
from .service import ServiceConnection, ServiceController
try:
from .operator import (
ModalJobOperatorBase,
ModalJobState,
active_job,
cancel_active_job,
close_active_job,
)
except ModuleNotFoundError as error:
if error.name != "bpy":
raise
__all__ = (
"JobCancelled",
"JobClient",
"JobContext",
"JobRequest",
"JobResult",
"JobServer",
"ServiceConnection",
"ServiceController",
"BlenderJobRuntime",
"JobOperatorBase",
"JobOperatorState",
"ServerConnection",
"ServerController",
)
if "ModalJobOperatorBase" in globals():
__all__ += (
"ModalJobOperatorBase",
"ModalJobState",
"active_job",
"cancel_active_job",
"close_active_job",
)

60
src/snapform/blendjob/client.py
View File

@@ -1,5 +1,6 @@
import json
from dataclasses import dataclass
from pathlib import Path
from urllib.error import HTTPError, URLError
from urllib.request import ProxyHandler, Request, build_opener
@@ -42,19 +43,13 @@ class JobClient:
def health(self):
return self.request("GET", "/health")
def submit(self, job_type, **parameters):
def submit(self, job_type, parameters):
return self.request(
"POST",
"/jobs",
{
"job_type": job_type,
"parameters": parameters,
},
{"job_type": str(job_type), "parameters": dict(parameters)},
)
def submit_request(self, request):
return self.request("POST", "/jobs", request.payload)
def status(self, job_id):
return self.request("GET", f"/jobs/{job_id}")
@@ -72,28 +67,35 @@ class JobClient:
def shutdown(self):
return self.request("POST", "/shutdown")
@dataclass(frozen=True, init=False)
class JobRequest:
"""Serializable request created before a Job is submitted."""
@dataclass(frozen=True)
class JobResult:
"""Successful Job value and files returned to a Blender Operator."""
job_type: str
parameters: dict
def __init__(self, job_type, parameters=None):
if isinstance(job_type, dict):
payload = job_type
job_type = payload.get("job_type", "")
parameters = payload.get("parameters", {})
object.__setattr__(self, "job_type", str(job_type))
object.__setattr__(self, "parameters", dict(parameters or {}))
job_id: str
directory: Path
value: object
status: dict
@classmethod
def create(cls, job_type, **parameters):
return cls(job_type, parameters)
def from_status(cls, status):
return cls(
job_id=str(status.get("job_id", "")),
directory=Path(status["directory"]),
value=status.get("result"),
status=dict(status),
)
@property
def payload(self):
return {
"job_type": self.job_type,
"parameters": self.parameters,
}
def file(self, name):
if not isinstance(self.value, dict):
raise RuntimeError("Job result does not contain named files")
try:
relative = self.value[name]
except KeyError:
raise RuntimeError(f"Job result does not contain file: {name}") from None
root = self.directory.resolve()
path = (root / str(relative)).resolve()
if path != root and root not in path.parents:
raise RuntimeError(f"Job result file escapes its directory: {relative}")
if not path.is_file():
raise RuntimeError(f"Job result file does not exist: {path}")
return path

61
src/snapform/blendjob/service.py → src/snapform/blendjob/controller.py
View File

@@ -4,18 +4,18 @@ import time
import uuid
from dataclasses import dataclass
from .client import JobClient
from .client import JobClient, JobResult
@dataclass
class ServiceConnection:
class ServerConnection:
port: int
pid: int
instance_id: str
client: JobClient
class ServiceController:
class ServerController:
"""Own a dedicated local JobServer process and its HTTP connection."""
def __init__(
@@ -61,7 +61,7 @@ class ServiceController:
self.log_file = None
raise
client = JobClient(self.host, port)
self.connection = ServiceConnection(
self.connection = ServerConnection(
port,
self.process.pid,
instance_id,
@@ -72,10 +72,10 @@ class ServiceController:
def health(self, timeout=0.3):
if self.connection is None:
raise RuntimeError("Service is not running")
raise RuntimeError("Server is not running")
health = self.connection.client.request("GET", "/health", timeout=timeout)
if health.get("instance_id") != self.connection.instance_id:
raise RuntimeError("Service instance does not match")
raise RuntimeError("Server instance does not match")
state = "BUSY" if health.get("busy") else "READY"
self.snapshot = {
**health,
@@ -97,7 +97,7 @@ class ServiceController:
if self.process is None:
self.start()
time.sleep(0.05)
raise RuntimeError(f"Service did not become ready; log: {self.log_path_factory()}")
raise RuntimeError(f"Server did not become ready; log: {self.log_path_factory()}")
def enable(self):
self.auto_start = True
@@ -129,8 +129,51 @@ class ServiceController:
self.auto_start = True
return self.start()
def submit(self, request):
return self.ensure().client.submit_request(request)
def submit(self, job_type, parameters):
return self.ensure().client.submit(job_type, parameters)
def request(
self,
job_type,
parameters,
*,
response=None,
progress=None,
cancel_check=None,
poll_interval=0.25,
):
"""Submit one Job and return its response after polling to completion."""
submitted = self.submit(job_type, parameters)
job_id = submitted["job_id"]
directory = submitted["directory"]
cancel_requested = False
try:
while True:
if (
not cancel_requested
and cancel_check is not None
and cancel_check()
):
self.cancel(job_id)
cancel_requested = True
status = self.status(job_id)
if progress is not None:
progress({"job_id": job_id, "directory": directory, **status})
state = status.get("state")
if state == "succeeded":
result = JobResult.from_status(
{"job_id": job_id, "directory": directory, **status}
)
if response is not None:
response(result)
return result
if state == "cancelled":
raise RuntimeError("Server request was cancelled")
if state == "failed":
raise RuntimeError(status.get("error") or "Server request failed")
time.sleep(max(float(poll_interval), 0.01))
finally:
self.mark_job_complete(job_id)
def status(self, job_id):
return self.ensure().client.status(job_id)

351
src/snapform/blendjob/environment.py Normal file
View File

@@ -0,0 +1,351 @@
import argparse
import hashlib
import json
import os
import platform
import queue
import shutil
import subprocess
import sys
import tarfile
import tempfile
import threading
import time
import urllib.request
import zipfile
from pathlib import Path
UV_VERSION = "0.11.3"
ENVIRONMENT_SCHEMA = 1
BASE_PACKAGES = (
"fastapi==0.139.2",
"uvicorn==0.51.0",
)
UV_ARTIFACTS = {
"x86_64-pc-windows-msvc": (
"uv-x86_64-pc-windows-msvc.zip",
"ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c",
),
"aarch64-apple-darwin": (
"uv-aarch64-apple-darwin.tar.gz",
"2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f",
),
"x86_64-unknown-linux-gnu": (
"uv-x86_64-unknown-linux-gnu.tar.gz",
"c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce",
),
}
UV_SOURCES = (
f"https://github.com/astral-sh/uv/releases/download/{UV_VERSION}",
f"https://mirror.omoolab.xyz/uv/{UV_VERSION}",
)
def normalized_environment(environment):
"""Validate and normalize a public Environment dictionary."""
if not isinstance(environment, dict):
raise TypeError("environment must be a dict")
python = str(environment.get("python", "")).strip()
if not python:
raise ValueError("environment['python'] is required")
packages = environment.get("packages", ())
platform_packages = environment.get("platform_packages", {})
if not isinstance(packages, (list, tuple)):
raise TypeError("environment['packages'] must be a list")
if not isinstance(platform_packages, dict):
raise TypeError("environment['platform_packages'] must be a dict")
normalized_platforms = {}
for name, values in platform_packages.items():
if not isinstance(values, (list, tuple)):
raise TypeError(f"platform_packages[{name!r}] must be a list")
normalized_platforms[str(name)] = [str(value) for value in values]
return {
"python": python,
"packages": [str(value) for value in packages],
"platform_packages": normalized_platforms,
}
def environment_digest(environment):
payload = json.dumps(
{
"schema": ENVIRONMENT_SCHEMA,
"base_packages": BASE_PACKAGES,
"environment": normalized_environment(environment),
},
sort_keys=True,
separators=(",", ":"),
).encode("utf-8")
return hashlib.sha256(payload).hexdigest()
def platform_name():
if sys.platform == "win32":
return "windows"
if sys.platform == "darwin":
return "macos"
return "linux"
def resolved_packages(environment):
environment = normalized_environment(environment)
platform_packages = environment["platform_packages"]
selected = platform_packages.get(
platform_name(),
platform_packages.get("default", ()),
)
return [*BASE_PACKAGES, *environment["packages"], *selected]
def environment_python(storage_root):
directory = Path(storage_root) / ".venv"
if os.name == "nt":
return directory / "Scripts" / "python.exe"
return directory / "bin" / "python"
def uv_path(storage_root):
filename = "uv.exe" if os.name == "nt" else "uv"
return Path(storage_root) / "tools" / filename
def uv_target():
machine = platform.machine().lower()
if sys.platform == "win32" and machine in {"amd64", "x86_64"}:
return "x86_64-pc-windows-msvc"
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
return "aarch64-apple-darwin"
if sys.platform.startswith("linux") and machine in {"amd64", "x86_64"}:
return "x86_64-unknown-linux-gnu"
raise RuntimeError(f"Unsupported uv platform: {sys.platform} {machine}")
def write_status(path, progress, message, *, stage, stages, error=None):
payload = {
"progress": min(max(float(progress), 0.0), 1.0),
"message": str(message),
"stage": int(stage),
"stages": int(stages),
"stage_label": "Environment",
}
if error:
payload["error"] = str(error)
Path(path).write_text(
json.dumps(payload, ensure_ascii=False),
encoding="utf-8",
)
def child_environment():
environment = os.environ.copy()
environment["PYTHONIOENCODING"] = "utf-8:backslashreplace"
environment["PYTHONUTF8"] = "1"
for name, value in urllib.request.getproxies().items():
environment.setdefault(f"{name.upper()}_PROXY", value)
return environment
def _download(url, destination, status, stage, stages, start, end):
request = urllib.request.Request(url, headers={"User-Agent": "BlendJob"})
with urllib.request.urlopen(request) as response, destination.open("wb") as output:
total = int(response.headers.get("Content-Length") or 0)
received = 0
while True:
chunk = response.read(256 * 1024)
if not chunk:
break
output.write(chunk)
received += len(chunk)
fraction = received / total if total else 0.5
progress = start + (end - start) * min(fraction, 1.0)
write_status(
status,
progress,
"Installing uv",
stage=stage,
stages=stages,
)
def _verify(path, checksum):
digest = hashlib.sha256()
with Path(path).open("rb") as stream:
for chunk in iter(lambda: stream.read(1024 * 1024), b""):
digest.update(chunk)
if digest.hexdigest().lower() != checksum.lower():
raise RuntimeError("uv archive checksum does not match")
def _extract_uv(archive, destination):
destination.parent.mkdir(parents=True, exist_ok=True)
with tempfile.TemporaryDirectory() as directory:
extracted = Path(directory)
if archive.name.endswith(".zip"):
with zipfile.ZipFile(archive) as package:
package.extractall(extracted)
else:
with tarfile.open(archive, "r:gz") as package:
package.extractall(extracted)
candidates = tuple(extracted.rglob(destination.name))
if not candidates:
raise RuntimeError("uv executable was not found in its archive")
shutil.copy2(candidates[0], destination)
if os.name != "nt":
destination.chmod(destination.stat().st_mode | 0o111)
def install_uv(storage_root, status, stage, stages, start=0.0, end=0.15):
destination = uv_path(storage_root)
if destination.is_file():
write_status(status, end, "uv is ready", stage=stage, stages=stages)
return destination
target = uv_target()
filename, checksum = UV_ARTIFACTS[target]
errors = []
destination.parent.mkdir(parents=True, exist_ok=True)
for source in UV_SOURCES:
descriptor, archive_name = tempfile.mkstemp(suffix=f"-{filename}")
os.close(descriptor)
archive = Path(archive_name)
try:
_download(
f"{source}/{filename}",
archive,
status,
stage,
stages,
start,
end * 0.8,
)
_verify(archive, checksum)
_extract_uv(archive, destination)
write_status(status, end, "uv is ready", stage=stage, stages=stages)
return destination
except Exception as error:
errors.append(f"{source}: {error}")
finally:
archive.unlink(missing_ok=True)
raise RuntimeError("Unable to install uv: " + "; ".join(errors))
def run_command(command, environment=None):
subprocess.check_call(
[str(value) for value in command],
env=environment or child_environment(),
)
def install_packages(
command,
status,
stage,
stages,
start=0.22,
end=0.98,
clock=time.monotonic,
):
process = subprocess.Popen(
[str(value) for value in command],
env=child_environment(),
stdout=subprocess.PIPE,
stderr=subprocess.STDOUT,
text=True,
encoding="utf-8",
errors="replace",
)
output = queue.Queue()
def read_output():
for line in process.stdout:
output.put(line)
output.put(None)
threading.Thread(target=read_output, daemon=True).start()
started = clock()
last_line = "Installing packages"
while True:
try:
line = output.get(timeout=0.2)
except queue.Empty:
line = ""
if line is None:
break
if line.strip():
last_line = line.strip()
print(last_line)
elapsed = clock() - started
fraction = min(elapsed / (elapsed + 30.0), 0.97)
write_status(
status,
start + (end - start) * fraction,
last_line,
stage=stage,
stages=stages,
)
return_code = process.wait()
if return_code:
raise subprocess.CalledProcessError(return_code, command)
def install(storage_root, environment, status, stages=1):
environment = normalized_environment(environment)
storage_root = Path(storage_root).resolve()
storage_root.mkdir(parents=True, exist_ok=True)
venv = storage_root / ".venv"
manifest = storage_root / "manifest.json"
expected_digest = environment_digest(environment)
uv = install_uv(storage_root, status, 1, stages)
manifest.unlink(missing_ok=True)
write_status(status, 0.16, "Creating Python Environment", stage=1, stages=stages)
run_command([uv, "venv", "--clear", "--python", environment["python"], venv])
write_status(status, 0.22, "Installing packages", stage=1, stages=stages)
packages = resolved_packages(environment)
install_packages(
[uv, "pip", "install", "--python", environment_python(storage_root), *packages],
status,
1,
stages,
)
manifest.write_text(
json.dumps(
{
"environment_hash": expected_digest,
"python": environment["python"],
"packages": packages,
},
indent=2,
),
encoding="utf-8",
)
write_status(status, 1.0, "Environment is ready", stage=1, stages=stages)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--storage-root", required=True)
parser.add_argument("--config", required=True)
parser.add_argument("--status", required=True)
parser.add_argument("--stages", type=int, default=1)
return parser.parse_args()
def main():
args = parse_args()
try:
environment = json.loads(Path(args.config).read_text(encoding="utf-8"))
install(args.storage_root, environment, args.status, args.stages)
except Exception as error:
write_status(
args.status,
1.0,
"Environment installation failed",
stage=1,
stages=args.stages,
error=error,
)
raise
if __name__ == "__main__":
main()

13
src/snapform/blendjob/installer/environment.py Normal file
View File

@@ -0,0 +1,13 @@
import sys
from pathlib import Path
BLENDJOB_DIRECTORY = Path(__file__).resolve().parent.parent
if str(BLENDJOB_DIRECTORY) not in sys.path:
sys.path.insert(0, str(BLENDJOB_DIRECTORY))
from environment import main
if __name__ == "__main__":
main()

260
src/snapform/blendjob/operator.py
View File

@@ -1,19 +1,18 @@
import json
import threading
from dataclasses import dataclass
from pathlib import Path
import bpy
from .client import JobRequest
from .client import JobResult
@dataclass
class ModalJobState:
"""Runtime state shared by all Blender modal Job operators."""
class JobOperatorState:
"""Internal state for one modal Job Operator invocation."""
runtime: object
operator: object
controller: object
data: object = None
cleanup: object = None
job_id: str = ""
directory: Path | None = None
started: bool = False
@@ -23,121 +22,110 @@ class ModalJobState:
progress_stage: int | None = None
progress: float = 0.0
def __getattr__(self, name):
if name == "result_directory":
return self.directory
if self.data is None:
raise AttributeError(name)
return getattr(self.data, name)
def operator_properties(operator):
"""Return every custom Blender Property declared by an Operator class."""
names = []
for class_type in reversed(type(operator).__mro__):
for name in getattr(class_type, "__annotations__", {}):
if name not in names:
names.append(name)
parameters = {
name: _json_value(getattr(operator, name))
for name in names
if hasattr(operator, name)
}
json.dumps(parameters)
return parameters
_active_job = None
def active_job():
return _active_job
def cancel_active_job():
job = active_job()
if job is None:
return
job.cancelled = True
if not job.started or not job.job_id:
return
def _json_value(value):
if value is None or isinstance(value, (bool, float, int, str)):
return value
if isinstance(value, (list, tuple)):
return [_json_value(item) for item in value]
try:
job.controller.cancel(job.job_id)
except RuntimeError:
pass
return [_json_value(item) for item in value]
except TypeError:
return value
def close_active_job():
global _active_job
if _active_job is not None:
if _active_job.job_id:
_active_job.controller.mark_job_complete(_active_job.job_id)
if _active_job.cleanup is not None:
_active_job.cleanup(_active_job.data)
_active_job.cleanup = None
_active_job = None
class ModalJobOperatorBase(bpy.types.Operator):
"""Ready-to-use Blender lifecycle for one HTTP Job request."""
class JobOperatorBase:
"""Blender mixin that turns one Operator into a complete remote Job."""
bl_options = {"INTERNAL"}
poll_interval = 0.25
starting_message = "Starting task"
controller_factory = None
job_runtime = None
job_type = ""
def get_service_controller(self):
"""Return the ServiceController used by this add-on."""
if self.controller_factory is None:
raise RuntimeError("Set controller_factory on the Job operator")
return self.controller_factory()
def __init_subclass__(cls, **kwargs):
super().__init_subclass__(**kwargs)
annotations = {}
for base in reversed(cls.__mro__[1:]):
annotations.update(getattr(base, "__annotations__", {}))
annotations.update(cls.__dict__.get("__annotations__", {}))
if annotations:
cls.__annotations__ = annotations
def create_job_request(self, context):
"""Create the JobRequest for this operator invocation."""
if not self.job_type:
raise RuntimeError("Set job_type on the Job operator")
return JobRequest(self.job_type, self.job_parameters(context))
def request(self, _context):
"""Return the complete JSON parameters submitted to the Server."""
return operator_properties(self)
def job_parameters(self, _context):
"""Return JSON-compatible business parameters for this invocation."""
return {}
def response(self, _context, _result):
"""Apply an optional successful result on Blender's main thread."""
def create_job_data(self, _context):
"""Return optional plug-in data retained until completion."""
return None
def cleanup(self):
"""Release invocation resources after every terminal outcome."""
def controller(self, runtime):
"""Return the operation controller used by this invocation."""
return runtime.server
def execute(self, context):
global _active_job
if active_job() is not None:
runtime = self._runtime()
if runtime.active_job is not None:
self.report({"WARNING"}, "Another task is running")
return {"CANCELLED"}
if not self.job_type:
self.report({"ERROR"}, "Set job_type on the Job Operator")
return {"CANCELLED"}
try:
request = self.create_job_request(context)
if not isinstance(request, JobRequest):
raise TypeError("create_job_request() must return JobRequest")
controller = self.get_service_controller()
data = self.create_job_data(context)
parameters = self.request(context)
if not isinstance(parameters, dict):
raise TypeError("request(context) must return a dict")
json.dumps(parameters)
controller = self.controller(runtime)
except (OSError, RuntimeError, TypeError, ValueError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
job = ModalJobState(
controller=controller,
data=data,
cleanup=self.cleanup_job_data,
)
_active_job = job
job = JobOperatorState(runtime, self, controller)
runtime.begin_job(job)
if hasattr(context.window_manager, "progress_begin"):
context.window_manager.progress_begin(0, 100)
self.update_job_ui(context, 0.0, self.get_starting_message(context))
runtime.update_ui(context, 0.0, self.starting_message)
self._timer = context.window_manager.event_timer_add(
self.poll_interval,
window=context.window,
)
context.window_manager.modal_handler_add(self)
self.redraw_job_ui(context, force=True)
runtime.redraw_ui(context, force=True)
job.start_thread = threading.Thread(
target=self._submit_job,
args=(job, request),
args=(job, parameters),
name="BlendJobSubmit",
daemon=True,
)
job.start_thread.start()
return {"RUNNING_MODAL"}
def _submit_job(self, job, request):
def _submit_job(self, job, parameters):
if job.cancelled:
job.started = True
return
try:
submitted = job.controller.submit(request)
submitted = job.controller.submit(self.job_type, parameters)
job.job_id = submitted["job_id"]
job.directory = Path(submitted["directory"])
job.started = True
@@ -147,53 +135,54 @@ class ModalJobOperatorBase(bpy.types.Operator):
job.start_error = str(error)
def modal(self, context, event):
runtime = self._runtime()
if event.type == "ESC":
cancel_active_job()
runtime.cancel_active()
if event.type != "TIMER":
return {"PASS_THROUGH"}
job = active_job()
job = runtime.active_job
if job is None:
self.remove_job_timer(context)
return {"CANCELLED"}
if job.start_error:
return self._finish_start_error(context, job)
self.report({"ERROR"}, f"Unable to start task: {job.start_error}")
return self._close(context, job, "Task failed to start", cancelled=True)
if not job.started:
return {"RUNNING_MODAL"}
if job.cancelled and not job.job_id:
return self._close(context, job, "Task cancelled", cancelled=True)
try:
status = job.controller.status(job.job_id)
except RuntimeError as error:
return self._finish_failure(context, job, str(error))
return self._finish_failure(context, job, error)
self.update_from_job_status(context, job, status)
state = status.get("state")
if state == "succeeded":
return self._finish_success(context, job, status)
if state in {"failed", "cancelled"}:
return self._finish_terminal(context, job, status)
if job.cancelled or state == "cancelled":
return self._close(context, job, "Task cancelled", cancelled=True)
return self._finish_failure(
context,
job,
status.get("error") or "Server task failed",
)
return {"RUNNING_MODAL"}
def _finish_start_error(self, context, job):
self.report({"ERROR"}, f"Unable to start task: {job.start_error}")
return self._close(context, job, "Task failed to start", cancelled=True)
def _finish_terminal(self, context, job, status):
if job.cancelled or status.get("state") == "cancelled":
return self._close(context, job, "Task cancelled", cancelled=True)
return self._finish_failure(
context,
job,
status.get("error") or "Service task failed",
)
def _finish_failure(self, context, job, error):
self.report({"ERROR"}, str(error))
return self._close(context, job, "Task failed", cancelled=True)
def _finish_success(self, context, job, status):
try:
message = self.on_job_succeeded(context, job, status)
except (OSError, RuntimeError, ValueError) as error:
complete_status = {
"job_id": job.job_id,
"directory": str(job.directory),
**status,
}
message = self.response(
context,
JobResult.from_status(complete_status),
)
except (OSError, RuntimeError, TypeError, ValueError) as error:
return self._finish_failure(context, job, error)
return self._close(
context,
@@ -202,18 +191,19 @@ class ModalJobOperatorBase(bpy.types.Operator):
cancelled=False,
)
def _finish_failure(self, context, job, error):
self.report({"ERROR"}, str(error))
return self._close(context, job, "Task failed", cancelled=True)
def _close(self, context, job, message, *, cancelled):
self.remove_job_timer(context)
if hasattr(context.window_manager, "progress_end"):
context.window_manager.progress_end()
close_active_job()
self.update_job_ui(context, 0.0 if cancelled else 1.0, message)
self.redraw_job_ui(context, force=True)
job.runtime.finish_job(job)
job.runtime.update_ui(context, 0.0 if cancelled else 1.0, message)
job.runtime.redraw_ui(context, force=True)
return {"CANCELLED"} if cancelled else {"FINISHED"}
def get_starting_message(self, _context):
return self.starting_message
def update_from_job_status(self, context, job, status):
reported = min(max(float(status.get("progress", job.progress)), 0.0), 1.0)
stage = status.get("stage")
@@ -223,30 +213,8 @@ class ModalJobOperatorBase(bpy.types.Operator):
else:
job.progress = max(job.progress, reported)
message = self.format_job_status(status, self.starting_message)
self.update_job_ui(context, job.progress, message)
self.redraw_job_ui(context, force=True)
def update_job_ui(self, context, progress, _message):
"""Update Blender's standard progress indicator by default."""
if hasattr(context.window_manager, "progress_update"):
context.window_manager.progress_update(int(progress * 100))
def on_job_succeeded(self, _context, _job, status):
"""Import results if needed and optionally return a final message."""
return status.get("message")
def cleanup_job_data(self, _data):
pass
def redraw_job_ui(self, _context, force=False):
del force
def remove_job_timer(self, context):
timer = getattr(self, "_timer", None)
if timer is None:
return
context.window_manager.event_timer_remove(timer)
self._timer = None
job.runtime.update_ui(context, job.progress, message)
job.runtime.redraw_ui(context, force=True)
def format_job_status(self, status, fallback):
message = status.get("message")
@@ -261,7 +229,27 @@ class ModalJobOperatorBase(bpy.types.Operator):
return f"{prefix} {stage_label}" if stage_label else prefix
return str(message or stage_label or fallback).strip()
def remove_job_timer(self, context):
timer = getattr(self, "_timer", None)
if timer is None:
return
context.window_manager.event_timer_remove(timer)
self._timer = None
def cancel(self, context):
cancel_active_job()
runtime = self._runtime()
job = runtime.active_job
runtime.cancel_active()
self.remove_job_timer(context)
self.redraw_job_ui(context, force=True)
if job is None:
return
if hasattr(context.window_manager, "progress_end"):
context.window_manager.progress_end()
runtime.finish_job(job)
runtime.update_ui(context, 0.0, "Task cancelled")
runtime.redraw_ui(context, force=True)
def _runtime(self):
if self.job_runtime is None:
raise RuntimeError("Use BlenderJobRuntime.JobOperatorBase")
return self.job_runtime

71
src/snapform/blendjob/runner.py Normal file
View File

@@ -0,0 +1,71 @@
import argparse
import importlib
import sys
import threading
from pathlib import Path
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--entrypoint", required=True)
parser.add_argument("--storage-root", required=True)
parser.add_argument("--host", default="127.0.0.1")
parser.add_argument("--port", type=int, required=True)
parser.add_argument("--instance-id", required=True)
parser.add_argument("--parent-pid", type=int, required=True)
return parser.parse_args()
def load_server(entrypoint):
path = Path(entrypoint).resolve()
if path.name == "__init__.py":
module_root = path.parent.parent
module_name = path.parent.name
else:
module_root = path.parent
module_name = path.stem
addon_root = Path(__file__).resolve().parent.parent
for directory in (module_root, addon_root):
if str(directory) not in sys.path:
sys.path.insert(0, str(directory))
module = importlib.import_module(module_name)
server = getattr(module, "server", None)
if server is None:
raise RuntimeError(f"Entrypoint does not export 'server': {entrypoint}")
return server
def main():
args = parse_args()
server = load_server(args.entrypoint)
server.bind(args.storage_root)
import uvicorn
from blendjob.server import watch_parent
app = server.create_app(args.instance_id)
config = uvicorn.Config(
app,
host=args.host,
port=args.port,
access_log=False,
log_level="warning",
)
http_server = uvicorn.Server(config)
def stop_http_server():
server.shutdown_event.wait()
http_server.should_exit = True
threading.Thread(target=stop_http_server, daemon=True).start()
threading.Thread(
target=watch_parent,
args=(args.parent_pid, server.shutdown_event),
name="BlendJobParentMonitor",
daemon=True,
).start()
http_server.run()
if __name__ == "__main__":
main()

596
src/snapform/blendjob/runtime.py Normal file
View File

@@ -0,0 +1,596 @@
import json
import os
import subprocess
import sys
import threading
from pathlib import Path
from .controller import ServerController
from .environment import environment_digest, normalized_environment
from .operator import JobOperatorBase
class EnvironmentController:
"""Install the declared Environment, then run arbitrary project setup."""
def __init__(self, runtime):
self.runtime = runtime
self.process = None
self.log_file = None
self.post_thread = None
self.post_status = {}
self.post_error = ""
self.cancelled = False
self._post_job_id = ""
self._post_stage = 1
self._lock = threading.Lock()
def submit(self, _job_type, _parameters):
runtime = self.runtime
runtime.storage_root()
self.post_thread = None
self.post_status = {}
self.post_error = ""
self.cancelled = False
self._post_job_id = ""
self._post_stage = 1
self.log_file = runtime.environment_log_path().open(
"w", encoding="utf-8"
)
environment = os.environ.copy()
environment["PYTHONIOENCODING"] = "utf-8:backslashreplace"
environment["PYTHONUTF8"] = "1"
self.process = subprocess.Popen(
runtime.install_command(),
stdout=self.log_file,
stderr=subprocess.STDOUT,
cwd=runtime.storage_root(),
creationflags=windows_creation_flags(),
env=environment,
)
return {
"job_id": "environment",
"directory": str(runtime.storage_root()),
}
def status(self, _job_id):
status = self._read_status()
if self.process is None or self.process.poll() is None:
return {
**status,
"state": "running",
"progress": status.get("progress", 0.0),
}
if self.process.returncode != 0:
return {
"message": "Environment installation failed",
**status,
"state": "failed",
"error": status.get("error", "Environment installer failed"),
}
if self.runtime.post_install is None:
return self._success_status()
if self.post_thread is None:
self.post_thread = threading.Thread(
target=self._run_post_install,
name="BlendJobPostInstall",
daemon=True,
)
self.post_thread.start()
if self.post_thread.is_alive():
with self._lock:
post_status = dict(self.post_status)
return {
"progress": post_status.get("progress", 0.0),
"message": post_status.get(
"message",
"Running post-install setup",
),
"stage": post_status.get("stage", 2),
"stage_label": post_status.get("stage_label", "Post Install"),
"state": "running",
}
return self._success_status()
def cancel(self, _job_id):
self.cancelled = True
if self.process is not None and self.process.poll() is None:
self.process.terminate()
if self._post_job_id:
try:
self.runtime.server.cancel(self._post_job_id)
except RuntimeError:
pass
return {"state": "cancelled"}
def mark_job_complete(self, _job_id):
if self.log_file is not None:
self.log_file.close()
self.log_file = None
def _run_post_install(self):
runtime = self.runtime
runtime._request_progress = self._update_post_progress
runtime._request_cancel_check = lambda: self.cancelled
try:
runtime.post_install(runtime)
except Exception as error:
self.post_error = str(error)
finally:
runtime._request_progress = None
runtime._request_cancel_check = None
def _update_post_progress(self, status):
job_id = str(status.get("job_id", ""))
with self._lock:
if job_id and job_id != self._post_job_id:
self._post_job_id = job_id
self._post_stage += 1
self.post_status = {
**status,
"stage": self._post_stage,
"stage_label": str(status.get("job_type", "Post Install")),
}
def _success_status(self):
if self.cancelled:
return {
"progress": 1.0,
"message": "Environment installation cancelled",
"state": "cancelled",
}
if self.post_error:
return {
"progress": 1.0,
"message": f"Environment is ready; post-install failed: {self.post_error}",
"state": "succeeded",
"post_install_error": self.post_error,
}
return {
"progress": 1.0,
"message": "Environment is ready",
"state": "succeeded",
}
def _read_status(self):
try:
return json.loads(
self.runtime.environment_status_path().read_text(
encoding="utf-8"
)
)
except (FileNotFoundError, json.JSONDecodeError, OSError):
return {}
class BlenderJobRuntime:
"""Own one add-on's Storage Root, Environment, Server and Blender UI."""
def __init__(
self,
server,
*,
storage_root,
environment,
post_install=None,
namespace=None,
):
self.server_definition = server
self.storage_root_factory = storage_root
self.environment = normalized_environment(environment)
self.environment_hash = environment_digest(self.environment)
if post_install is not None and not callable(post_install):
raise TypeError("post_install must be callable")
self.post_install = post_install
self.namespace = namespace or _default_namespace(server.name)
self.active_job = None
self.progress = 0.0
self.message = "Ready"
self._operator_base = None
self._operator_classes = None
self._timer_registered = False
self._registered = False
self._status_bar_draw = self._create_status_bar_draw()
self._request_progress = None
self._request_cancel_check = None
self._environment_controller = EnvironmentController(self)
self.server = ServerController(
self._server_command,
cwd_factory=self.storage_root,
log_path_factory=self.server_log_path,
creationflags=windows_creation_flags(),
)
self.JobOperatorBase = self.operator_base()
def storage_root(self):
configured = (
self.storage_root_factory()
if callable(self.storage_root_factory)
else self.storage_root_factory
)
root = Path(configured).expanduser().resolve()
root.mkdir(parents=True, exist_ok=True)
return root
def environment_directory(self):
return self.storage_root() / ".venv"
def environment_python(self):
directory = self.environment_directory()
if os.name == "nt":
return directory / "Scripts" / "python.exe"
return directory / "bin" / "python"
def environment_manifest_path(self):
return self.storage_root() / "manifest.json"
def environment_config_path(self):
return self.storage_root() / "environment.json"
def environment_status_path(self):
return self.storage_root() / "environment-status.json"
def environment_log_path(self):
return self.storage_root() / "environment.log"
def server_log_path(self):
return self.storage_root() / "server.log"
def environment_manifest(self):
try:
value = json.loads(
self.environment_manifest_path().read_text(encoding="utf-8")
)
except (FileNotFoundError, json.JSONDecodeError, OSError):
return {}
return value if isinstance(value, dict) else {}
def environment_ready(self):
return (
self.environment_python().is_file()
and self.environment_manifest().get("environment_hash")
== self.environment_hash
)
def write_environment_config(self):
path = self.environment_config_path()
path.write_text(
json.dumps(self.environment, indent=2),
encoding="utf-8",
)
return path
def install_command(self):
installer = Path(__file__).with_name("installer") / "environment.py"
config = self.write_environment_config()
return [
sys.executable,
str(installer),
"--storage-root",
str(self.storage_root()),
"--config",
str(config),
"--status",
str(self.environment_status_path()),
"--stages",
"1",
]
def request(self, job_type, parameters, response=None):
"""Directly request one Server Job and return its JobResult."""
if self._request_cancel_check is not None and self._request_cancel_check():
raise RuntimeError("Server request was cancelled")
return self.server.request(
job_type,
parameters,
response=response,
progress=self._request_progress,
cancel_check=self._request_cancel_check,
)
def operator_base(self):
if self._operator_base is not None:
return self._operator_base
runtime = self
class BoundJobOperator(JobOperatorBase):
job_runtime = runtime
BoundJobOperator.__name__ = "JobOperatorBase"
self._operator_base = BoundJobOperator
return BoundJobOperator
def operator_classes(self):
if self._operator_classes is not None:
return self._operator_classes
import bpy
runtime = self
operator_base = self.JobOperatorBase
class InstallEnvironment(operator_base, bpy.types.Operator):
bl_idname = f"{runtime.namespace}.install_environment"
bl_label = "Install Environment"
job_type = "install-environment"
starting_message = "Starting Environment installation"
def request(self, _context):
if not bpy.app.online_access:
raise RuntimeError(
"Online Access is disabled in Blender Preferences"
)
return {}
def controller(self, _runtime):
runtime.server.stop()
return runtime._environment_controller
class CancelJob(bpy.types.Operator):
bl_idname = f"{runtime.namespace}.cancel_job"
bl_label = "Cancel Job"
@classmethod
def poll(cls, _context):
return runtime.active_job is not None
def execute(self, _context):
runtime.cancel_active()
return {"FINISHED"}
class StartServer(bpy.types.Operator):
bl_idname = f"{runtime.namespace}.start_server"
bl_label = "Start Server"
def execute(self, _context):
try:
runtime.server.enable()
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class StopServer(bpy.types.Operator):
bl_idname = f"{runtime.namespace}.stop_server"
bl_label = "Stop Server"
def execute(self, _context):
runtime.server.stop()
runtime.redraw_ui()
return {"FINISHED"}
class RestartServer(bpy.types.Operator):
bl_idname = f"{runtime.namespace}.restart_server"
bl_label = "Restart Server"
def execute(self, _context):
try:
runtime.server.restart()
except (OSError, RuntimeError) as error:
self.report({"ERROR"}, str(error))
return {"CANCELLED"}
return {"FINISHED"}
class OpenServerLog(bpy.types.Operator):
bl_idname = f"{runtime.namespace}.open_server_log"
bl_label = "Open Server Log"
def execute(self, _context):
path = runtime.server_log_path()
target = path if path.is_file() else path.parent
result = bpy.ops.wm.path_open(filepath=str(target))
if "FINISHED" not in result:
self.report({"ERROR"}, "Unable to open the Server log")
return {"CANCELLED"}
return {"FINISHED"}
self._operator_classes = (
InstallEnvironment,
CancelJob,
StartServer,
StopServer,
RestartServer,
OpenServerLog,
)
return self._operator_classes
def register(self):
if self._registered:
return
import bpy
for class_type in self.operator_classes():
bpy.utils.register_class(class_type)
bpy.types.STATUSBAR_HT_header.append(self._status_bar_draw)
self._registered = True
self.enable()
def unregister(self):
if not self._registered:
return
import bpy
self.disable()
bpy.types.STATUSBAR_HT_header.remove(self._status_bar_draw)
for class_type in reversed(self.operator_classes()):
bpy.utils.unregister_class(class_type)
self._registered = False
def begin_job(self, job):
if self.active_job is not None:
raise RuntimeError("Another task is running")
self.active_job = job
def finish_job(self, job):
if job.job_id:
job.controller.mark_job_complete(job.job_id)
try:
job.operator.cleanup()
finally:
if self.active_job is job:
self.active_job = None
def cancel_active(self):
job = self.active_job
if job is None:
return
job.cancelled = True
if not job.started or not job.job_id:
return
try:
job.controller.cancel(job.job_id)
except RuntimeError:
pass
def close_active(self):
job = self.active_job
if job is not None:
self.finish_job(job)
def update_ui(self, context, progress, message):
self.progress = min(max(float(progress), 0.0), 1.0)
self.message = str(message)
window_manager = getattr(context, "window_manager", None)
if window_manager is not None and hasattr(
window_manager, "progress_update"
):
window_manager.progress_update(int(self.progress * 100))
self.redraw_ui(context)
def redraw_ui(self, context=None, force=False):
try:
import bpy
except ModuleNotFoundError:
return
context = context or bpy.context
window_manager = getattr(context, "window_manager", None)
for window in getattr(window_manager, "windows", ()):
screen = getattr(window, "screen", None)
for area in getattr(screen, "areas", ()):
if area.type in {"STATUSBAR", "VIEW_3D", "PREFERENCES"}:
area.tag_redraw()
for region in getattr(area, "regions", ()):
region.tag_redraw()
if not force:
return
workspace = getattr(context, "workspace", None)
if workspace is not None:
try:
workspace.status_text_set_internal(None)
except (AttributeError, RuntimeError, TypeError):
pass
def server_status(self):
return dict(self.server.snapshot)
def server_busy(self):
return self.active_job is not None or self.server.snapshot.get(
"state"
) == "BUSY"
def resource(self, name):
return self.server.resource(name)
def clear_resource(self, name):
return self.server.clear_resource(name)
def enable(self):
try:
import bpy
except ModuleNotFoundError:
return
timers = getattr(bpy.app, "timers", None)
if timers is None or timers.is_registered(self._poll_server):
return
timers.register(self._poll_server, first_interval=1.0, persistent=True)
self._timer_registered = True
def disable(self):
self.cancel_active()
self.close_active()
try:
import bpy
except ModuleNotFoundError:
bpy = None
timers = getattr(getattr(bpy, "app", None), "timers", None)
if timers is not None and timers.is_registered(self._poll_server):
timers.unregister(self._poll_server)
self._timer_registered = False
self.server.detach()
def _poll_server(self):
interval = self.server.poll(available=self.environment_ready())
self.redraw_ui()
return interval
def _draw_status_bar(self, owner, _context):
if self.active_job is None:
return
row = owner.layout.row(align=True)
progress_row = row.row(align=True)
progress_row.ui_units_x = 18.0
progress_row.progress(
factor=self.progress,
type="BAR",
text=self.message,
)
row.operator(
f"{self.namespace}.cancel_job",
text="",
icon="X",
)
def _create_status_bar_draw(self):
runtime = self
def draw_status_bar(owner, context):
runtime._draw_status_bar(owner, context)
return draw_status_bar
def _server_command(self, port, instance_id):
entrypoint_path = self._server_entrypoint()
if entrypoint_path is None:
raise RuntimeError("Job Server does not export a Server entrypoint")
runner = Path(__file__).with_name("runner.py")
return [
str(self.environment_python()),
"-u",
str(runner),
"--entrypoint",
str(entrypoint_path),
"--storage-root",
str(self.storage_root()),
"--host",
"127.0.0.1",
"--port",
str(port),
"--instance-id",
instance_id,
"--parent-pid",
str(os.getpid()),
]
def _server_entrypoint(self):
candidates = []
for module in tuple(sys.modules.values()):
if getattr(module, "server", None) is not self.server_definition:
continue
module_path = getattr(module, "__file__", None)
if module_path:
candidates.append(Path(module_path).resolve())
package_entries = [path for path in candidates if path.name == "__init__.py"]
if package_entries:
return min(package_entries, key=lambda path: len(path.parts))
if candidates:
return min(candidates, key=lambda path: len(path.parts))
return self.server_definition.entrypoint_path
def _default_namespace(server_name):
return str(server_name).split()[0].lower().replace("_", "-")
def windows_creation_flags():
if os.name != "nt":
return 0
return subprocess.CREATE_NO_WINDOW

172
src/snapform/blendjob/server.py
View File

@@ -1,8 +1,10 @@
import os
import sys
import threading
import time
import traceback
import uuid
from concurrent.futures import ThreadPoolExecutor
from contextlib import asynccontextmanager
from pathlib import Path
@@ -17,10 +19,19 @@ class JobCancelled(RuntimeError):
class JobContext:
"""Mutable state and cancellation API passed to a registered job handler."""
def __init__(self, job_id, job_type, parameters, resources, directory):
def __init__(
self,
job_id,
job_type,
parameters,
resources,
storage_root,
directory,
):
self.job_id = job_id
self.job_type = job_type
self.parameters = parameters
self.storage_root = Path(storage_root)
self.directory = Path(directory)
self._resources = resources
self._lock = threading.Lock()
@@ -34,8 +45,7 @@ class JobContext:
"directory": str(self.directory),
}
def update(self, progress, message, error=None, state=None, **details):
"""Publish a progress snapshot returned by ``GET /jobs/{job_id}``."""
def _update(self, progress, message, error=None, state=None, **details):
if state is None:
if error:
state = "failed"
@@ -57,6 +67,15 @@ class JobContext:
with self._lock:
self._status = snapshot
def progress(self, progress, message, **details):
"""Publish running progress from a Job handler."""
self._update(progress, message, state="running", **details)
def succeed(self, result=None, message="Task complete"):
"""Publish a successful terminal result."""
details = {"result": result} if result is not None else {}
self._update(1.0, message, state="succeeded", **details)
def snapshot(self):
with self._lock:
return dict(self._status)
@@ -86,26 +105,59 @@ class JobContext:
try:
return self._resources[name]
except KeyError:
raise KeyError(f"Unknown service resource: {name}") from None
raise KeyError(f"Unknown Server Resource: {name}") from None
class JobServer:
"""Small FastAPI wrapper with decorator-based job registration."""
def __init__(self, name, *, jobs_directory=None, max_job_history=100):
def __init__(
self,
name,
*,
storage_root=None,
max_job_history=100,
):
self.name = name
self.max_job_history = max(int(max_job_history), 0)
self.handlers = {}
self.resources = {}
self.resource_factories = {}
self.jobs = {}
self.futures = {}
self.active_job_id = None
self.lock = threading.Lock()
self.executor = ThreadPoolExecutor(
max_workers=1,
thread_name_prefix="BlendJobWorker",
)
self.shutdown_event = threading.Event()
self.started_at = time.time()
self.jobs_directory = Path(jobs_directory or Path.cwd() / "jobs")
self.storage_root = Path(storage_root) if storage_root is not None else None
self.jobs_directory = (
self.storage_root / "jobs" if self.storage_root is not None else None
)
self.entrypoint_module = None
self.entrypoint_path = None
def bind(self, storage_root):
"""Bind the Server to the Runtime-owned Storage Root."""
storage_root = Path(storage_root).resolve()
if self.storage_root is not None:
current = self.storage_root.resolve()
if current != storage_root:
raise RuntimeError(
f"Job Server is already bound to another Storage Root: {current}"
)
self.storage_root = storage_root
self.jobs_directory = storage_root / "jobs"
for name, factory in tuple(self.resource_factories.items()):
if name not in self.resources:
self.resources[name] = factory(self)
return self
def add_resource(self, name, resource):
"""Register shared state owned for the full service lifetime."""
"""Register shared state owned for the full Server lifetime."""
if not name or "/" in name:
raise ValueError(f"Invalid resource name: {name}")
with self.lock:
@@ -114,6 +166,20 @@ class JobServer:
self.resources[name] = resource
return resource
def resource(self, name):
"""Register a Resource factory initialized after Storage Root binding."""
def register(factory):
if not name or "/" in name:
raise ValueError(f"Invalid resource name: {name}")
if name in self.resources or name in self.resource_factories:
raise ValueError(f"Resource is already registered: {name}")
self.resource_factories[name] = factory
if self.storage_root is not None:
self.resources[name] = factory(self)
return factory
return register
def resource_snapshot(self, name):
try:
resource = self.resources[name]
@@ -129,9 +195,9 @@ class JobServer:
}
def clear_resource(self, name):
"""Clear an idle service resource without racing a new job submission."""
"""Clear an idle Server Resource without racing a new Job submission."""
with self.lock:
if self.active_job_id is not None:
if self.active_job_id is not None or self._queued_job_count():
return False
try:
resource = self.resources[name]
@@ -144,6 +210,7 @@ class JobServer:
return True
def close(self):
self.executor.shutdown(wait=False, cancel_futures=True)
for resource in reversed(tuple(self.resources.values())):
close = getattr(resource, "close", None)
if close is not None:
@@ -155,6 +222,12 @@ class JobServer:
if job_type in self.handlers:
raise ValueError(f"Job type is already registered: {job_type}")
self.handlers[job_type] = handler
if self.entrypoint_module is None:
self.entrypoint_module = handler.__module__
module = sys.modules.get(handler.__module__)
module_path = getattr(module, "__file__", None)
if module_path:
self.entrypoint_path = Path(module_path).resolve()
return handler
return register
@@ -162,26 +235,28 @@ class JobServer:
def snapshot(self, instance_id):
with self.lock:
active = self.jobs.get(self.active_job_id)
queued_jobs = self._queued_job_count()
result = {
"ready": True,
"service": self.name,
"server": self.name,
"instance_id": instance_id,
"busy": active is not None,
"busy": active is not None or queued_jobs > 0,
"active_job": active.snapshot() if active else None,
"queued_jobs": queued_jobs,
"started_at": self.started_at,
"resources": self.resource_snapshots(),
}
return result
def submit(self, job_type, parameters, *, job_id=None):
if self.storage_root is None or self.jobs_directory is None:
raise RuntimeError("Bind the Job Server to a Storage Root before submitting")
job_id = job_id or uuid.uuid4().hex
if job_type not in self.handlers:
raise KeyError(job_type)
if not job_id or Path(job_id).name != job_id or job_id in {".", ".."}:
raise ValueError(f"Invalid job id: {job_id}")
with self.lock:
if self.active_job_id is not None:
return None
if job_id in self.jobs:
raise ValueError(f"Job id already exists: {job_id}")
self.jobs_directory.mkdir(parents=True, exist_ok=True)
@@ -195,13 +270,33 @@ class JobServer:
job_type,
parameters,
self.resources,
self.storage_root,
directory,
)
self.jobs[job_id] = context
self.active_job_id = job_id
self._prune_jobs()
self.futures[job_id] = self.executor.submit(self.run, context)
return context
def cancel(self, job_id):
with self.lock:
context = self.jobs.get(job_id)
future = self.futures.get(job_id)
if context is None:
return None
if future is not None and future.cancel():
context.request_cancel()
context._update(1.0, "Task cancelled", state="cancelled")
self.futures.pop(job_id, None)
return context
return context if context.request_cancel() else None
def _queued_job_count(self):
return sum(
context.snapshot()["state"] == "queued"
for context in self.jobs.values()
)
def _prune_jobs(self):
completed = [
identifier
@@ -214,25 +309,42 @@ class JobServer:
self.jobs.pop(identifier, None)
def run(self, context):
context.update(0.0, "Service accepted the task", state="running")
with self.lock:
if context.snapshot()["state"] in TERMINAL_STATES:
return
self.active_job_id = context.job_id
context._update(0.0, "Server accepted the task", state="running")
try:
context.check_cancelled()
self.handlers[context.job_type](context, context.parameters)
result = self.handlers[context.job_type](
context,
context.parameters,
)
context.check_cancelled()
if context.snapshot()["state"] not in TERMINAL_STATES:
context.update(1.0, "Task complete", state="succeeded")
snapshot = context.snapshot()
if snapshot["state"] not in {"failed", "cancelled"}:
context.succeed(
result,
message=snapshot.get("message") or "Task complete",
)
except JobCancelled:
context.update(1.0, "Task cancelled", state="cancelled")
context._update(1.0, "Task cancelled", state="cancelled")
except Exception as error:
traceback.print_exc()
context.update(1.0, "Task failed", error=error, state="failed")
context._update(1.0, "Task failed", error=error, state="failed")
finally:
with self.lock:
if self.active_job_id == context.job_id:
self.active_job_id = None
self.futures.pop(context.job_id, None)
def create_app(self, instance_id):
from fastapi import BackgroundTasks, FastAPI, HTTPException
def create_app(self, instance_id, storage_root=None):
from fastapi import FastAPI, HTTPException
if storage_root is not None:
self.bind(storage_root)
elif self.storage_root is None:
raise RuntimeError("Storage Root is required to create the Job Server app")
@asynccontextmanager
async def lifespan(_app):
@@ -276,14 +388,11 @@ class JobServer:
except TypeError as error:
raise HTTPException(status_code=405, detail=str(error))
if not cleared:
raise HTTPException(status_code=409, detail="Service is busy")
raise HTTPException(status_code=409, detail="Server is busy")
return {"resource": resource_name, "state": "cleared"}
@app.post("/jobs", status_code=202)
def submit_job(
payload: dict,
background_tasks: BackgroundTasks,
):
def submit_job(payload: dict):
job_type = str(payload.get("job_type", payload.get("command", "")))
parameters = payload.get("parameters")
if not job_type or not isinstance(parameters, dict):
@@ -294,9 +403,6 @@ class JobServer:
raise HTTPException(status_code=422, detail="Unknown job type")
except ValueError as error:
raise HTTPException(status_code=409, detail=str(error))
if context is None:
raise HTTPException(status_code=409, detail="Service is busy")
background_tasks.add_task(self.run, context)
return context.snapshot()
@app.get("/jobs/{job_id}")
@@ -308,8 +414,8 @@ class JobServer:
@app.delete("/jobs/{job_id}", status_code=202)
def cancel_job(job_id: str):
context = self.jobs.get(job_id)
if context is None or not context.request_cancel():
context = self.cancel(job_id)
if context is None:
raise HTTPException(status_code=404, detail="Active job was not found")
return context.snapshot()
@@ -322,7 +428,7 @@ class JobServer:
def watch_parent(parent_pid, shutdown_event, interval=1.0):
"""Stop a dedicated service after its owning Blender process exits."""
"""Stop a dedicated Server after its owning Blender process exits."""
if not parent_pid:
return
while not shutdown_event.wait(interval):

521
src/snapform/client.py
View File

@@ -1,521 +0,0 @@
import json
import os
import platform
import subprocess
import sys
import tempfile
import threading
from pathlib import Path
import bpy
from .blendjob import JobRequest, ServiceController
from .env_revision import MANIFEST_FILENAME, MANIFEST_REVISION
from .models import (
BEN2_MODEL,
MOGE2_MODEL,
get_downloadable_model,
model_directory_name,
selected_model_key,
)
SERVICE_ADDRESS = "127.0.0.1"
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
_service_controller = None
_environment_controller = None
class EnvironmentController:
"""Adapt the local Environment bootstrap to the modal controller protocol."""
def __init__(self):
self.process = None
self.log_file = None
def submit(self, _request):
stop_service()
paths = install_paths()
self.log_file = paths["log"].open("w", encoding="utf-8")
environment = os.environ.copy()
environment["PYTHONIOENCODING"] = "utf-8:backslashreplace"
environment["PYTHONUTF8"] = "1"
self.process = subprocess.Popen(
install_env_command(paths["status"]),
stdout=self.log_file,
stderr=subprocess.STDOUT,
cwd=storage_root(),
creationflags=creation_flags(),
env=environment,
)
return {"job_id": "environment", "directory": str(storage_root())}
def status(self, _job_id):
status = read_bootstrap_status(install_paths()["status"]) or {}
if self.process is None or self.process.poll() is None:
return {**status, "state": "running", "progress": status.get("progress", 0.0)}
if self.process.returncode == 0:
return {"progress": 1.0, "message": "Environment is ready", **status,
"state": "succeeded"}
return {"message": "Environment installation failed", **status,
"state": "failed", "error": status.get("error", "Bootstrap failed")}
def cancel(self, _job_id):
if self.process is not None and self.process.poll() is None:
self.process.terminate()
return {"state": "cancelled"}
def mark_job_complete(self, _job_id):
if self.log_file is not None:
self.log_file.close()
self.log_file = None
def environment_controller():
global _environment_controller
if _environment_controller is None:
_environment_controller = EnvironmentController()
return _environment_controller
def preferences():
addon = bpy.context.preferences.addons.get(__package__)
return addon.preferences if addon is not None else None
def storage_root():
addon_preferences = preferences()
configured = getattr(addon_preferences, "storage_root", "")
root = (
Path(bpy.path.abspath(configured)).expanduser()
if configured
else DEFAULT_STORAGE_ROOT
)
root.mkdir(parents=True, exist_ok=True)
return root
def venv_directory():
return storage_root() / ".venv"
def venv_python():
directory = venv_directory()
if os.name == "nt":
return directory / "Scripts" / "python.exe"
return directory / "bin" / "python"
def env_installed():
ready_path = storage_root() / MANIFEST_FILENAME
if not venv_python().is_file() or not ready_path.is_file():
return False
try:
manifest = json.loads(ready_path.read_text(encoding="utf-8"))
except (json.JSONDecodeError, OSError):
return False
return (
isinstance(manifest, dict)
and manifest.get("manifest_revision") == MANIFEST_REVISION
)
def models_directory():
path = storage_root() / "models"
path.mkdir(parents=True, exist_ok=True)
return path
def model_directory(model_key):
root = models_directory()
directory = root / model_directory_name(model_key)
old_moge_directory = root / "moge2_base_normal"
if (
model_key == "MOGE2_VITB_NORMAL"
and not directory.is_dir()
and old_moge_directory.is_dir()
):
old_moge_directory.rename(directory)
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):
model = get_downloadable_model(model_key)
path = model_directory(model_key)
return all(
any(path.glob(pattern))
for pattern in model.ready_patterns
)
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:
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 before environment installation."""
devices = ["cpu"]
if sys.platform == "win32":
devices.insert(0, "directml")
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 env_devices() or machine_devices()
DEVICE_LABELS = {
"cuda": ("CUDA", "NVIDIA CUDA"),
"directml": ("DirectML", "Windows GPU (DirectML)"),
"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 temporary_input_path():
directory = Path(tempfile.mkdtemp(prefix="snapform-input-"))
return directory / "input.png"
def temporary_input_directory():
return Path(tempfile.mkdtemp(prefix="snapform-input-"))
def install_paths():
"""Return shared bootstrap paths outside the BlendJob workspace."""
root = storage_root()
return {
"job_id": "",
"log": root / "environment.log",
"status": root / "environment-status.json",
"result": root,
}
def uv_executable():
"""Return the downloaded uv path inside the shared storage."""
return str(uv_cache_path())
def uv_cache_path():
if os.name == "nt":
return storage_root() / "tools" / "uv.exe"
return storage_root() / "tools" / "uv"
def install_env_command(status_path):
executable = uv_executable()
script = Path(__file__).with_name("install_env.py")
command = [
sys.executable,
str(script),
"--uv",
executable,
"--venv-dir",
str(venv_directory()),
"--manifest",
str(storage_root() / MANIFEST_FILENAME),
"--models-dir",
str(models_directory()),
"--status",
str(status_path),
]
return command
def service_log_path():
return storage_root() / "service.log"
def _service_command(port, instance_id):
return [
str(venv_python()),
"-u",
str(Path(__file__).with_name("ai_service") / "server.py"),
"--host",
SERVICE_ADDRESS,
"--port",
str(port),
"--instance-id",
instance_id,
"--parent-pid",
str(os.getpid()),
]
def service_controller():
global _service_controller
if _service_controller is None:
_service_controller = ServiceController(
_service_command,
cwd_factory=storage_root,
log_path_factory=service_log_path,
host=SERVICE_ADDRESS,
creationflags=creation_flags(),
)
return _service_controller
def start_service():
return service_controller().start()
def ensure_service(timeout=60.0):
connection = service_controller().ensure(timeout)
_redraw_service_ui()
return connection
def stop_service():
service_controller().stop()
def restart_service():
return service_controller().restart()
def enable_service():
return service_controller().enable()
def clear_service_models():
return ensure_service().client.clear_resource("models")
def detach_service_client():
service_controller().detach()
def schedule_service_start():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_timer):
return
timers.register(_service_timer, first_interval=1.0, persistent=True)
def unschedule_service():
timers = getattr(bpy.app, "timers", None)
if timers is None:
return
if timers.is_registered(_service_timer):
timers.unregister(_service_timer)
def service_status():
return dict(service_controller().snapshot)
def service_busy():
return service_controller().snapshot.get("state") == "BUSY"
def _service_timer():
interval = service_controller().poll(available=env_installed())
_redraw_service_ui()
return interval
def _redraw_service_ui():
if threading.current_thread() is not threading.main_thread():
return
try:
from .status_bar import redraw_ui
redraw_ui(bpy.context, force=True)
except (AttributeError, RuntimeError):
pass
def download_model(model_key):
get_downloadable_model(model_key)
return JobRequest("download-model", {"model": model_key})
def remove_background(settings, input_path, model_key, *, start_frame=1, device=None):
return _background_request(
"remove-background", settings, input_path, model_key, start_frame, device
)
def debug_ben2(settings, input_path, model_key, *, start_frame=1, device=None):
return _background_request(
"debug-ben2", settings, input_path, model_key, start_frame, device
)
def _background_request(job_type, settings, input_path, model_key, start_frame, device):
return JobRequest(job_type, {
"input": str(input_path),
"model": model_key,
"device": (device or settings.device).lower(),
"video_frames": getattr(settings, "video_frames", 1),
"start_frame": start_frame,
})
def extract_frames(settings, input_path, *, start_frame=1):
return JobRequest("extract-frames", {
"input": str(input_path),
"video_frames": settings.video_frames,
"start_frame": start_frame,
})
def generate_da3_depth(settings, input_path, model_key, *, process_res, use_ray_pose,
remove_background=False, start_frame=1, device=None):
return JobRequest("generate-da3-depth", {
**_da3_common(settings, input_path, model_key, process_res, device, start_frame),
"export_format": "mini_npz",
"use_ray_pose": 1 if use_ray_pose else 0,
"inference_mode": "per_frame",
"remove_background_model": BEN2_MODEL.key if remove_background else "",
})
def generate_da3_surface(settings, input_path, model_key, *, process_res, use_ray_pose,
remove_background=False, start_frame=1, device=None):
return JobRequest("generate-da3-surface", {
**_da3_common(settings, input_path, model_key, process_res, device, start_frame),
"export_format": "mesh_data",
"use_ray_pose": 1 if use_ray_pose else 0,
"inference_mode": "per_frame",
"remove_background_model": BEN2_MODEL.key if remove_background else "",
"mesh_filter_edges": bool(getattr(settings, "mesh_filter_edges", True)),
"mesh_max_edge_length": float(getattr(settings, "mesh_max_edge_length", 0.1)),
})
def debug_da3(settings, input_path=None, model_key=None, *, device=None):
export_format = settings.export_format
selected_key = model_key or selected_model_key(settings, export_format)
return JobRequest("debug-da3", {
**_da3_common(
settings,
input_path or Path(bpy.path.abspath(settings.input_path)),
selected_key,
settings.process_res,
device,
1,
),
"export_format": export_format,
"use_ray_pose": 1 if settings.use_ray_pose else 0,
"inference_mode": "joint",
"remove_background_model": "",
"mesh_filter_edges": bool(getattr(settings, "mesh_filter_edges", True)),
"mesh_max_edge_length": float(getattr(settings, "mesh_max_edge_length", 0.1)),
})
def _da3_common(settings, input_path, model_key, process_res, device, start_frame):
return {
"input": str(input_path),
"model": model_key,
"process_res": process_res,
"process_res_method": "upper_bound_resize",
"conf_thresh_percentile": settings.conf_thresh_percentile,
"num_max_points": settings.num_max_points,
"show_cameras": 1 if settings.show_cameras else 0,
"gs_views_interval": settings.gs_views_interval,
"video_frames": settings.video_frames,
"start_frame": start_frame,
"device": (device or settings.device).lower(),
"ref_view_strategy": settings.ref_view_strategy,
}
def generate_moge2_depth(settings, input_path, model_key=MOGE2_MODEL.key, *,
resolution_level=5, remove_background=False,
start_frame=1, device=None):
return JobRequest("generate-moge2-depth", {
**_moge2_common(settings, input_path, model_key, resolution_level, device, start_frame),
"remove_background_model": BEN2_MODEL.key if remove_background else "",
})
def generate_moge2_surface(settings, input_path, model_key=MOGE2_MODEL.key, *,
resolution_level=5, remove_background=False,
start_frame=1, device=None):
return JobRequest("generate-moge2-surface", {
**_moge2_common(settings, input_path, model_key, resolution_level, device, start_frame),
"remove_background_model": BEN2_MODEL.key if remove_background else "",
})
def debug_moge2(settings, input_path=None, model_key=None, *, device=None):
return JobRequest("debug-moge2", {
**_moge2_common(
settings,
input_path or Path(bpy.path.abspath(settings.input_path)),
model_key or settings.moge_model,
settings.moge_resolution_level,
device,
1,
),
"export_format": settings.moge_export_format,
"include_points": True,
})
def _moge2_common(settings, input_path, model_key, resolution_level, device, start_frame):
return {
"input": str(input_path),
"model": model_key,
"resolution_level": int(resolution_level),
"video_frames": settings.video_frames,
"start_frame": start_frame,
"device": (device or settings.device).lower(),
}
def creation_flags():
if os.name != "nt":
return 0
return subprocess.CREATE_NO_WINDOW
def read_bootstrap_status(path):
try:
return json.loads(Path(path).read_text(encoding="utf-8"))
except (FileNotFoundError, json.JSONDecodeError, OSError):
return None

4
src/snapform/env_revision.py
View File

@@ -1,4 +0,0 @@
"""Shared environment metadata for the Blender UI and the install subprocess."""
MANIFEST_FILENAME = "manifest.json"
MANIFEST_REVISION = 8

921
src/snapform/install_env.py
View File

@@ -1,921 +0,0 @@
import argparse
import hashlib
import http.client
import json
import os
import platform
import queue
import shutil
import subprocess
import sys
import tarfile
import tempfile
import threading
import time
import urllib.request
import zipfile
from pathlib import Path
SCRIPT_DIRECTORY = Path(__file__).resolve().parent
if str(SCRIPT_DIRECTORY) not in sys.path:
sys.path.insert(0, str(SCRIPT_DIRECTORY))
from env_revision import MANIFEST_REVISION
from models import BEN2_MODEL, DOWNLOADABLE_MODELS, MOGE2_MODEL
PYPI_MIRRORS = (
"https://pypi.tuna.tsinghua.edu.cn/simple",
"https://mirrors.aliyun.com/pypi/simple/",
)
PYPI_DEFAULT_INDEX = "https://pypi.org/simple/"
FASTAPI_VERSION = "0.139.2"
UVICORN_VERSION = "0.51.0"
ONNXRUNTIME_VERSION = "1.24.4"
NUMPY_VERSION = "2.4.2"
PILLOW_VERSION = "12.1.1"
AV_VERSION = "17.0.0"
TRIMESH_VERSION = "4.11.2"
HUGGINGFACE_HUB_VERSION = "1.4.1"
UV_VERSION = "0.11.3"
UV_OFFICIAL_CONNECT_TIMEOUT = 3.0
UV_MIRROR_CONNECT_TIMEOUT = 15.0
DOWNLOAD_CHUNK_SIZE = 256 * 1024
UV_OFFICIAL_BASE_URL = (
f"https://github.com/astral-sh/uv/releases/download/{UV_VERSION}"
)
UV_MIRROR_BASE_URL = f"https://mirror.omoolab.xyz/uv/{UV_VERSION}"
UV_ARTIFACTS = {
"x86_64-pc-windows-msvc": (
"uv-x86_64-pc-windows-msvc.zip",
"ae681c0aaec7cc96af184648cb88d73f8393ed60fa5880abdd6bdb910f9b227c",
),
"aarch64-apple-darwin": (
"uv-aarch64-apple-darwin.tar.gz",
"2bc3d0c7bf2bd08325b1e170abac6f7e5b3346e1d4eab3370d17cefec934996f",
),
"x86_64-unknown-linux-gnu": (
"uv-x86_64-unknown-linux-gnu.tar.gz",
"c0f3236f146e55472663cfbcc9be3042a9f1092275bbe3fe2a56a6cbfd3da5ce",
),
}
_resolved_pypi_index = None
DEFAULT_MODEL_KEYS = (MOGE2_MODEL.key, BEN2_MODEL.key)
RAMP_RATE = 0.02
CREEP_CAP = 0.60
CHECK_CODE = (
"import json; import fastapi, uvicorn, onnxruntime; "
"import av, numpy, PIL, trimesh; "
"print('SNAPFORM_DEVICES=' + json.dumps({"
"'onnxruntime': onnxruntime.__version__, "
"'providers': onnxruntime.get_available_providers()}))"
)
class _ConnectHTTPSConnection(http.client.HTTPSConnection):
"""Limit TCP/TLS setup without limiting artifact reads."""
connect_timeout = UV_OFFICIAL_CONNECT_TIMEOUT
def connect(self):
try:
self.timeout = self.connect_timeout
super().connect()
finally:
self.timeout = None
self.sock.settimeout(None)
class _MirrorConnectHTTPSConnection(_ConnectHTTPSConnection):
connect_timeout = UV_MIRROR_CONNECT_TIMEOUT
class _ConnectHTTPSHandler(urllib.request.HTTPSHandler):
def __init__(self, connection_type):
super().__init__()
self.connection_type = connection_type
def https_open(self, request):
return self.do_open(self.connection_type, request)
def build_uv_download_opener(connection_type):
return urllib.request.build_opener(
urllib.request.ProxyHandler(),
_ConnectHTTPSHandler(connection_type),
)
_uv_download_openers = {
UV_OFFICIAL_BASE_URL: build_uv_download_opener(_ConnectHTTPSConnection),
UV_MIRROR_BASE_URL: build_uv_download_opener(
_MirrorConnectHTTPSConnection
),
}
def configure_standard_streams():
"""Windows 图形进程可能继承 GBK stdout，安装日志统一写为 UTF-8。"""
for stream in (sys.stdout, sys.stderr):
reconfigure = getattr(stream, "reconfigure", None)
if callable(reconfigure):
reconfigure(encoding="utf-8", errors="backslashreplace")
def write_status(
path,
progress,
message,
error=None,
stage=None,
stages=None,
stage_label=None,
):
payload = {"progress": progress, "message": message}
if error:
payload["error"] = error
if stage is not None:
payload["stage"] = stage
payload["stages"] = stages
payload["stage_label"] = stage_label
Path(path).write_text(
json.dumps(payload, ensure_ascii=False),
encoding="utf-8",
)
def child_environment(environment=None):
"""Inherit the current environment and inject system/registry proxy settings into child processes (uv only reads proxy env variables)."""
process_environment = os.environ.copy()
process_environment["PYTHONIOENCODING"] = "utf-8:backslashreplace"
process_environment["PYTHONUTF8"] = "1"
for key, value in urllib.request.getproxies().items():
env_key = f"{key.upper()}_PROXY"
if env_key not in process_environment:
process_environment[env_key] = value
if _resolved_pypi_index:
process_environment["UV_INDEX_URL"] = _resolved_pypi_index
if environment:
process_environment.update(environment)
return process_environment
def run(command, environment=None):
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
subprocess.check_call(
[str(value) for value in command],
env=child_environment(environment),
)
def parse_package_count(line):
try:
return int(line.split()[1])
except (IndexError, ValueError):
return 0
def resolve_indexes(args):
"""确定用户指定的 PyPI 源，优先命令行参数，其次 uv/pip 的环境变量。"""
global _resolved_pypi_index
_resolved_pypi_index = (
getattr(args, "pypi_index", None)
or os.environ.get("UV_INDEX_URL")
or os.environ.get("PIP_INDEX_URL")
)
def _run_install_once(
status_path,
stage,
stages,
label,
command,
environment,
start,
end,
clock,
):
"""Run a uv install command once and map package-count events onto the stage progress range."""
print("Running:", subprocess.list2cmdline([str(value) for value in command]))
process_environment = child_environment(environment)
process_environment["UV_LOG"] = "debug"
process = subprocess.Popen(
[str(value) for value in command],
env=process_environment,
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()
def eased(fraction):
return start + min(max(fraction, 0.0), 1.0) * (end - start)
installed = []
total = 0
displayed = start
target = start
started = clock()
last_tick = started
last_heartbeat = 0.0
def detail(message):
write_status(
status_path,
displayed,
message,
stage=stage,
stages=stages,
stage_label=label,
)
while True:
try:
raw_line = output_queue.get(timeout=0.2)
empty = False
except queue.Empty:
raw_line = None
empty = True
now = clock()
elapsed = now - started
dt = now - last_tick
last_tick = now
if total:
effective = max(target, eased(CREEP_CAP))
if effective > displayed:
displayed = min(effective, displayed + dt * RAMP_RATE)
if raw_line is None and not empty:
break
if empty:
if elapsed - last_heartbeat >= 5.0:
last_heartbeat = elapsed
# uv 的下载阶段没有增量进度，只报告已解析的包总数，
# 让用户对规模有预期，而不承诺剩余时间或已耗时。
message = (
f"Downloading {total} packages"
if total
else "Downloading packages"
)
write_status(
status_path,
displayed,
message,
stage=stage,
stages=stages,
stage_label=label,
)
continue
line = raw_line.rstrip()
print(line)
if line.startswith("Resolved ") and "package" in line:
total = parse_package_count(line)
resolved_progress = eased(0.02)
displayed = max(displayed, resolved_progress)
target = max(target, resolved_progress)
detail(f"Resolving {total} packages")
elif line.startswith("Prepared ") and "package" in line:
done = parse_package_count(line)
if total:
target = max(target, eased(0.10 + 0.60 * done / total))
detail(f"Downloading packages {done}/{total}")
elif line.startswith("Installed ") and "package" in line:
done = parse_package_count(line)
if total:
target = max(target, eased(0.75 + 0.25 * done / total))
detail(f"Installing packages {done}/{total}")
elif line.startswith("+"):
entry = line[1:].strip()
installed.append(entry)
return_code = process.wait()
if return_code != 0:
raise subprocess.CalledProcessError(return_code, command)
return installed
def run_install(
status_path,
stage,
stages,
label,
command,
environment=None,
start=0.0,
end=1.0,
clock=time.monotonic,
):
"""Run a uv install command, retrying across PyPI indexes on failure."""
if _resolved_pypi_index is None:
# 未指定源时优先国内镜像，pypi.org 作为最终兜底。
indexes = (*PYPI_MIRRORS, PYPI_DEFAULT_INDEX)
elif _resolved_pypi_index == PYPI_DEFAULT_INDEX:
# 明确选择官方源时，失败后回退国内镜像。
indexes = (PYPI_DEFAULT_INDEX, *PYPI_MIRRORS)
else:
# 用户指定了特定源，不再回退。
indexes = (_resolved_pypi_index,)
attempts = [
dict(environment or {}, UV_INDEX_URL=index) for index in indexes
]
for position, attempt_environment in enumerate(attempts):
try:
return _run_install_once(
status_path,
stage,
stages,
label,
command,
attempt_environment,
start,
end,
clock,
)
except subprocess.CalledProcessError:
if position + 1 == len(attempts):
raise
print(
"Retrying with next PyPI index: "
f"{attempts[position + 1]['UV_INDEX_URL']}"
)
def venv_python(venv_dir):
if sys.platform == "win32":
return Path(venv_dir) / "Scripts" / "python.exe"
return Path(venv_dir) / "bin" / "python"
def install_stages(args, venv_dir, python, uv):
"""Build the lightweight ONNX-only runtime and model download stages."""
runtime_package = (
f"onnxruntime-directml=={ONNXRUNTIME_VERSION}"
if sys.platform == "win32"
else f"onnxruntime=={ONNXRUNTIME_VERSION}"
)
stages = [
(
"Environment",
[uv, "venv", "--clear", "--python", "3.12", venv_dir],
None,
"plain",
),
(
"ONNX Runtime",
[
uv,
"pip",
"install",
"-v",
"--python",
python,
f"fastapi=={FASTAPI_VERSION}",
f"huggingface-hub=={HUGGINGFACE_HUB_VERSION}",
f"numpy=={NUMPY_VERSION}",
f"av=={AV_VERSION}",
runtime_package,
f"pillow=={PILLOW_VERSION}",
f"trimesh=={TRIMESH_VERSION}",
f"uvicorn=={UVICORN_VERSION}",
],
None,
"install",
),
]
models = DOWNLOADABLE_MODELS
worker_script = Path(__file__).with_name("ai_service") / "jobs.py"
models_dir = Path(args.models_dir)
for model_key in DEFAULT_MODEL_KEYS:
model = models[model_key]
download_command = [
python,
worker_script,
"download",
"--repository",
model.repository,
"--model-dir",
models_dir / model.directory_name,
"--status",
args.status,
"--label",
model.label,
]
for pattern in model.download_patterns:
download_command.extend(["--allow-pattern", pattern])
for pattern in model.ready_patterns:
download_command.extend(["--required-pattern", pattern])
if model.r2_directory:
download_command.extend(
["--r2-directory", model.r2_directory]
)
for filename, size, checksum in model.r2_files:
download_command.extend(
["--r2-file", f"{filename}|{size}|{checksum}"]
)
stages.append(
(
model.label,
download_command,
None,
"download",
)
)
return stages
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 (
isinstance(payload, dict)
and payload.get("manifest_revision") == MANIFEST_REVISION
)
def installed_devices():
devices = []
machine = platform.machine().lower()
if sys.platform == "darwin" and machine in {"arm64", "aarch64"}:
devices.append("mps")
elif sys.platform == "win32":
devices.append("directml")
devices.append("cpu")
return devices
def verify_command(python):
return [python, "-c", CHECK_CODE]
def run_verify(
python,
status_path,
start=0.0,
end=1.0,
clock=time.monotonic,
stage=None,
stages=None,
stage_label=None,
):
"""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(
[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,
"Checking Environment",
stage=stage,
stages=stages,
stage_label=stage_label,
)
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":
if machine in {"arm64", "aarch64"}:
return "aarch64-pc-windows-msvc", "zip"
return "x86_64-pc-windows-msvc", "zip"
if sys.platform == "darwin":
arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64"
return f"{arch}-apple-darwin", "tar.gz"
if sys.platform.startswith("linux"):
arch = "aarch64" if machine in {"arm64", "aarch64"} else "x86_64"
return f"{arch}-unknown-linux-gnu", "tar.gz"
raise RuntimeError(f"Unsupported platform for uv download: {sys.platform}")
def archive_is_valid(path):
if zipfile.is_zipfile(path):
return True
try:
with tarfile.open(path) as archive:
archive.getmembers()
return True
except (OSError, tarfile.TarError):
return False
def extract_uv(archive_path, uv_path):
executable_name = "uv.exe" if sys.platform == "win32" else "uv"
with tempfile.TemporaryDirectory() as directory:
if zipfile.is_zipfile(archive_path):
with zipfile.ZipFile(archive_path) as archive:
archive.extractall(directory)
else:
with tarfile.open(archive_path) as archive:
archive.extractall(directory)
source = next(Path(directory).rglob(executable_name))
shutil.copy2(source, uv_path)
if sys.platform != "win32":
os.chmod(uv_path, 0o755)
def file_sha256(path):
digest = hashlib.sha256()
with Path(path).open("rb") as stream:
while chunk := stream.read(1024 * 1024):
digest.update(chunk)
return digest.hexdigest()
def format_download_rate(rate):
size = float(rate)
units = ("B/s", "KB/s", "MB/s", "GB/s")
unit_index = 0
while size >= 1024 and unit_index < len(units) - 1:
size /= 1024
unit_index += 1
precision = 0 if unit_index == 0 else 1
return f"{size:.{precision}f} {units[unit_index]}"
def download_uv_archive(
destination,
status_path,
start,
end,
stage_fields,
clock=time.monotonic,
):
target, _archive_type = uv_download_target()
artifact = UV_ARTIFACTS.get(target)
if artifact is None:
raise RuntimeError(f"SnapForm does not provide uv for {target}")
filename, expected_checksum = artifact
failures = []
sources = (
(UV_OFFICIAL_BASE_URL, ""),
(UV_MIRROR_BASE_URL, " (mirror)"),
)
for base_url, source_label in sources:
destination.unlink(missing_ok=True)
url = f"{base_url}/{filename}"
write_status(
status_path,
start,
f"Downloading uv{source_label}",
**stage_fields,
)
request = urllib.request.Request(
url,
headers={"User-Agent": "SnapForm uv bootstrap"},
)
try:
with _uv_download_openers[base_url].open(request) as response:
total = int(response.headers.get("Content-Length", 0))
downloaded = 0
started_at = clock()
last_report_at = started_at
with destination.open("wb") as output:
while chunk := response.read(DOWNLOAD_CHUNK_SIZE):
output.write(chunk)
downloaded += len(chunk)
now = clock()
if now - last_report_at < 0.20 and downloaded < total:
continue
last_report_at = now
elapsed = max(now - started_at, 1e-6)
rate = downloaded / elapsed
message = (
f"Downloading uv{source_label} "
f"{format_download_rate(rate)}"
)
progress = start
if total:
progress += (end - start) * downloaded / total
write_status(
status_path,
min(progress, end),
message,
**stage_fields,
)
actual_checksum = file_sha256(destination)
if actual_checksum != expected_checksum:
raise RuntimeError(
f"SHA-256 mismatch: {actual_checksum}"
)
if not archive_is_valid(destination):
raise RuntimeError("Downloaded uv archive is invalid")
return destination
except Exception as error:
destination.unlink(missing_ok=True)
failures.append(f"{url}: {error}")
raise RuntimeError("Unable to download uv; " + "; ".join(failures))
def install_uv(
uv,
status_path,
stage=None,
stages=None,
stage_label=None,
start=0.0,
end=1.0,
):
"""Download, verify, and extract the pinned uv binary."""
stage_fields = {
"stage": stage,
"stages": stages,
"stage_label": stage_label,
}
Path(uv).parent.mkdir(parents=True, exist_ok=True)
archive_path = Path(uv).with_name("uv-download.tmp")
try:
download_uv_archive(
archive_path,
status_path,
start,
end,
stage_fields,
)
write_status(
status_path,
end,
"Extracting uv",
**stage_fields,
)
extract_uv(archive_path, uv)
finally:
archive_path.unlink(missing_ok=True)
def run_stage(
status_path,
stage,
stages,
label,
command,
environment,
kind,
start=0.0,
end=1.0,
):
if stage is None:
opening_message = f"Downloading {label}"
elif kind == "install":
opening_message = "Preparing packages"
else:
opening_message = ""
write_status(
status_path,
start,
opening_message,
stage=stage,
stages=stages,
stage_label=label,
)
installed = []
if kind == "install":
installed = run_install(
status_path,
stage,
stages,
label,
command,
environment,
start=start,
end=end,
)
elif kind == "download":
progress_command = [
*command,
"--progress-start",
str(start),
"--progress-end",
str(end),
]
if stage is not None:
progress_command.extend(
[
"--stage",
str(stage),
"--stages",
str(stages),
"--stage-label",
label,
]
)
run(progress_command, environment)
else:
run(command, environment)
write_status(
status_path,
end,
f"{label} ready" if stage is None else "",
stage=stage,
stages=stages,
stage_label=label,
)
return installed
def install(args):
resolve_indexes(args)
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 = venv_python(venv_dir)
stages = install_stages(args, venv_dir, python, uv)
env_stages = [stage for stage in stages if stage[3] != "download"]
model_stages = [stage for stage in stages if stage[3] == "download"]
total_stages = int(uv_missing) + 1 + len(model_stages)
stage = 1
installed_packages = []
if uv_missing:
install_uv(
uv,
args.status,
stage=stage,
stages=total_stages,
stage_label="uv",
start=0.0,
end=1.0,
)
stage += 1
environment_label = "ONNX Environment"
for index, (_label, command, environment, kind) in enumerate(env_stages):
if index >= 2:
raise RuntimeError("Unexpected ONNX Environment stage")
stage_start, stage_end = ((0.0, 0.05) if index == 0 else (0.05, 1.0))
stage_installed = run_stage(
args.status,
stage,
total_stages,
environment_label,
command,
environment,
kind,
start=stage_start,
end=stage_end,
)
if isinstance(stage_installed, list):
installed_packages.extend(stage_installed)
stage += 1
devices = installed_devices()
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(
{
"manifest_revision": MANIFEST_REVISION,
"python": "3.12",
"onnxruntime": ONNXRUNTIME_VERSION,
"fastapi": FASTAPI_VERSION,
"uvicorn": UVICORN_VERSION,
"runtime_variant": (
"directml" if sys.platform == "win32" else "default"
),
"devices": devices,
"packages": sorted(set(installed_packages)),
},
indent=2,
),
encoding="utf-8",
)
failed_models = []
for label, command, environment, kind in model_stages:
try:
run_stage(
args.status,
stage,
total_stages,
label,
command,
environment,
kind,
start=0.0,
end=1.0,
)
except Exception as error:
print(f"Model download failed: {label}: {error}")
failed_models.append(label)
write_status(
args.status,
1.0,
f"Skipping {label}; continuing",
stage=stage,
stages=total_stages,
stage_label=label,
)
stage += 1
if failed_models:
labels = ", ".join(failed_models)
write_status(
args.status,
1.0,
f"Environment is installed; retry {labels} from Models",
stage=total_stages,
stages=total_stages,
stage_label=model_stages[-1][0],
)
return
write_status(
args.status,
1.0,
"Environment and recommended models ready",
stage=total_stages,
stages=total_stages,
stage_label=model_stages[-1][0] if model_stages else environment_label,
)
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--uv", 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("--status", required=True)
parser.add_argument("--pypi-index", default=None)
return parser.parse_args()
def main():
configure_standard_streams()
args = parse_args()
try:
install(args)
except Exception as error:
write_status(
args.status,
1.0,
"Environment installation failed",
str(error),
)
raise
if __name__ == "__main__":
main()

54
src/snapform/job_runtime.py Normal file
View File

@@ -0,0 +1,54 @@
from pathlib import Path
import bpy
from .blendjob import BlenderJobRuntime
from .server import server
DEFAULT_STORAGE_ROOT = Path.home() / ".snapform"
def configured_storage_root():
addon = bpy.context.preferences.addons.get(__package__)
configured = getattr(getattr(addon, "preferences", None), "storage_root", "")
if configured:
return Path(bpy.path.abspath(configured))
return DEFAULT_STORAGE_ROOT
ENVIRONMENT = {
"python": "3.12",
"packages": [
"numpy==2.4.2",
"pillow==12.1.1",
"av==17.0.0",
"trimesh==4.11.2",
"huggingface-hub==1.4.1",
],
"platform_packages": {
"windows": ["onnxruntime-directml==1.24.4"],
"default": ["onnxruntime==1.24.4"],
},
}
def post_install(job_runtime):
errors = []
for model in ("MOGE2_VITS_NORMAL", "BEN2_BASE"):
try:
job_runtime.request("download-model", {"model": model})
except RuntimeError as error:
errors.append(str(error))
if errors:
raise RuntimeError("; ".join(errors))
runtime = BlenderJobRuntime(
server,
storage_root=configured_storage_root,
environment=ENVIRONMENT,
post_install=post_install,
)
JobOperatorBase = runtime.JobOperatorBase

1277
src/snapform/operators.py
View File

File diff suppressed because it is too large Load Diff

85
src/snapform/panel.py
View File

@@ -1,14 +1,16 @@
import bpy
from .models import BEN2_MODEL, MODELS, MOGE2_MODELS, selected_model_key
from .operators import active_job
from . import client
from .client import model_ready, env_installed
from .job_runtime import runtime
from .properties import BEN2_MODEL_KEY, model_catalog, model_ready
class AiServicePanel(bpy.types.Panel):
bl_label = "AI Service"
bl_idname = "SNAPFORM_PT_service"
def selected_model_key(settings):
return settings.gs_model if settings.export_format == "gs_ply" else settings.model
class ServerPanel(bpy.types.Panel):
bl_label = "Job Server"
bl_idname = "SNAPFORM_PT_server"
bl_space_type = "VIEW_3D"
bl_region_type = "UI"
bl_category = "SnapForm"
@@ -16,8 +18,8 @@ class AiServicePanel(bpy.types.Panel):
def draw(self, _context):
layout = self.layout
running = active_job() is not None or client.service_busy()
if not env_installed():
running = runtime.server_busy()
if not runtime.environment_ready():
layout.label(
text="Environment is not installed",
icon="ERROR",
@@ -31,7 +33,7 @@ class AiServicePanel(bpy.types.Panel):
)
return
status = client.service_status()
status = runtime.server_status()
state = status.get("state", "STOPPED")
icons = {
"READY": "RADIOBUT_ON",
@@ -54,41 +56,37 @@ class AiServicePanel(bpy.types.Panel):
if state == "READY":
controls = row.row(align=True)
operator = controls.operator(
"snapform.control_service",
"snapform.restart_server",
text="",
icon="FILE_REFRESH",
)
operator.action = "RESTART"
operator = controls.operator(
"snapform.control_service",
"snapform.stop_server",
text="",
icon="PAUSE",
)
operator.action = "STOP"
elif state == "STOPPED":
operator = row.operator(
"snapform.control_service",
row.operator(
"snapform.start_server",
text="Start",
icon="PLAY",
)
operator.action = "START"
else:
row.label(
text=status.get("message", state.title()),
icon=icons.get(state, "QUESTION"),
)
if state == "ERROR":
operator = row.operator(
"snapform.control_service",
row.operator(
"snapform.start_server",
text="Start",
icon="PLAY",
)
operator.action = "START"
if state in {"READY", "BUSY"}:
cached = (
status.get("resources", {})
.get("models", {})
.get("model_manager", {})
.get("loaded", {})
)
cached_names = [name for name in cached.values() if name]
@@ -102,12 +100,11 @@ class AiServicePanel(bpy.types.Panel):
)
clear = cached_row.row(align=True)
clear.enabled = state == "READY"
operator = clear.operator(
"snapform.control_service",
clear.operator(
"snapform.clear_models",
text="",
icon="TRASH",
)
operator.action = "CLEAR"
class ModelsPanel(bpy.types.Panel):
@@ -121,34 +118,41 @@ class ModelsPanel(bpy.types.Panel):
@classmethod
def poll(cls, _context):
return env_installed()
return runtime.environment_ready()
def draw(self, _context):
layout = self.layout
running = active_job() is not None or client.service_busy()
running = runtime.server_busy()
for model in MOGE2_MODELS.values():
self._draw_model(layout, model, running)
for model in MODELS.values():
if model.enabled:
models = model_catalog()
for model in models:
if model.get("family") == "moge2":
self._draw_model(layout, model, running)
for model in models:
if model.get("family") == "da3" and model.get("enabled"):
self._draw_model(layout, model, running)
layout.separator(type="LINE")
self._draw_model(layout, BEN2_MODEL, running)
ben2 = next(
(model for model in models if model.get("key") == BEN2_MODEL_KEY),
None,
)
if ben2:
self._draw_model(layout, ben2, running)
def _draw_model(self, layout, model, running):
row = layout.row(align=True)
row.label(text=model.label)
if model_ready(model.key):
row.label(text=model["label"])
if model.get("ready"):
row.label(text="Ready", icon="CHECKMARK")
else:
download_row = row.row(align=True)
download_row.enabled = env_installed() and not running
download_row.enabled = runtime.environment_ready() and not running
operator = download_row.operator(
"snapform.download_model",
text="Download",
icon="IMPORT",
)
operator.model_key = model.key
operator.model = model["key"]
class DebugPanel(bpy.types.Panel):
@@ -228,8 +232,7 @@ class Da3OutputDebugPanel(bpy.types.Panel):
mesh_box.prop(settings, "mesh_max_edge_length")
export_row = layout.row()
export_row.enabled = (
active_job() is None
and not client.service_busy()
not runtime.server_busy()
and model_ready(selected_model_key(settings))
)
export_row.operator(
@@ -258,8 +261,7 @@ class Moge2ApiDebugPanel(bpy.types.Panel):
layout.prop(settings, "moge_resolution_level")
run_row = layout.row()
run_row.enabled = (
active_job() is None
and not client.service_busy()
not runtime.server_busy()
and model_ready(settings.moge_model)
)
run_row.operator(
@@ -286,9 +288,8 @@ class Ben2RemovalDebugPanel(bpy.types.Panel):
layout.prop(settings, "background_input_path")
background_row = layout.row()
background_row.enabled = (
active_job() is None
and not client.service_busy()
and model_ready(BEN2_MODEL.key)
not runtime.server_busy()
and model_ready(BEN2_MODEL_KEY)
)
background_row.operator(
"snapform.debug_ben2",

21
src/snapform/preferences.py
View File

@@ -1,8 +1,7 @@
from pathlib import Path
import bpy
from .client import device_enum_items
from .job_runtime import DEFAULT_STORAGE_ROOT, runtime
from .properties import device_enum_items
class SnapFormPreferences(bpy.types.AddonPreferences):
@@ -15,7 +14,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
"all Blender versions can use the same directory"
),
subtype="DIR_PATH",
default=str(Path.home() / ".snapform"),
default=str(DEFAULT_STORAGE_ROOT),
)
device: bpy.props.EnumProperty(
name="Device",
@@ -28,19 +27,16 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
)
def draw(self, _context):
from .operators import active_job
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"Environment: {venv_directory()}")
if env_installed():
layout.label(text=f"Environment: {runtime.environment_directory()}")
if runtime.environment_ready():
layout.label(text="Environment is installed", icon="CHECKMARK")
else:
install_row = layout.row()
install_row.enabled = active_job() is None
install_row.enabled = runtime.active_job is None
install_row.operator(
"snapform.install_environment",
text="Install Environment",
@@ -48,8 +44,7 @@ class SnapFormPreferences(bpy.types.AddonPreferences):
)
layout.separator()
operator = layout.operator(
"snapform.control_service",
text="Open Service Log",
"snapform.open_server_log",
text="Open Server Log",
icon="TEXT",
)
operator.action = "LOG"

152
src/snapform/properties.py
View File

@@ -1,7 +1,142 @@
import platform
import sys
import bpy
from .models import gs_model_enum_items, model_enum_items, moge2_model_enum_items
from .client import device_enum_items
from .job_runtime import runtime
BEN2_MODEL_KEY = "BEN2_BASE"
DEFAULT_DA3_MODEL_KEY = "BASE"
DEFAULT_MOGE2_MODEL_KEY = "MOGE2_VITS_NORMAL"
DEVICE_LABELS = {
"cuda": ("CUDA", "NVIDIA CUDA"),
"directml": ("DirectML", "Windows GPU (DirectML)"),
"mps": ("MPS", "Apple Metal (MPS)"),
"cpu": ("CPU", "CPU"),
}
_models_cache = {}
def supported_devices():
devices = runtime.environment_manifest().get("devices")
if isinstance(devices, list) and devices:
return [str(device).lower() for device in devices]
if sys.platform == "win32":
return ["directml", "cpu"]
if sys.platform == "darwin" and platform.machine().lower() in {
"arm64",
"aarch64",
}:
return ["mps", "cpu"]
return ["cpu"]
def device_enum_items(_owner, _context):
return [
(name.upper(), *DEVICE_LABELS[name], index)
for index, name in enumerate(supported_devices())
]
def models_status(refresh=False):
global _models_cache
status = runtime.server_status().get("resources", {}).get("model_manager", {})
if status:
_models_cache = status
if status and not refresh:
return status
if _models_cache and not refresh:
return _models_cache
if not runtime.environment_ready():
return {}
try:
status = runtime.resource("model_manager")
_models_cache = status
return status
except RuntimeError:
return _models_cache
def model_ready(model_key):
return any(
model.get("key") == model_key and model.get("ready")
for model in models_status().get("catalog", ())
)
def model_catalog():
return models_status().get("catalog", ())
def model_record(model_key):
return next(
(model for model in model_catalog() if model.get("key") == model_key),
None,
)
def _model_items(models):
return [
(
model["key"],
model["label"] if model.get("ready") else f"{model['label']} (Missing)",
f"{model['description']}; license: {model['license']}",
index,
)
for index, model in enumerate(models)
]
def model_enum_items(_owner, _context):
return _model_items(
model
for model in model_catalog()
if model.get("family") == "da3" and model.get("enabled")
)
def gs_model_enum_items(_owner, _context):
return _model_items(
model
for model in model_catalog()
if model.get("family") == "da3"
and model.get("enabled")
and model.get("supports_gs")
)
def moge2_model_enum_items(_owner, _context):
return _model_items(
model for model in model_catalog() if model.get("family") == "moge2"
)
def depth_plane_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") in {"moge2", "da3"} and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
items = _model_items(models)
items.append(
("ORIGINAL", "Original", "Use the original image as depth", len(items))
)
return items
def surface_model_enum_items(_owner, _context):
models = [
model
for model in model_catalog()
if model.get("family") in {"moge2", "da3"}
and model.get("enabled")
]
models.sort(key=lambda model: model.get("key") != DEFAULT_MOGE2_MODEL_KEY)
return _model_items(models)
class SnapFormSettings(bpy.types.PropertyGroup):
@@ -185,16 +320,3 @@ class SnapFormSettings(bpy.types.PropertyGroup):
default=1,
min=1,
)
status: bpy.props.StringProperty(
name="Status",
default="Ready",
options={"SKIP_SAVE"},
)
progress: bpy.props.FloatProperty(
name="Progress",
default=0.0,
min=0.0,
max=1.0,
subtype="FACTOR",
options={"SKIP_SAVE"},
)

64
src/snapform/server/__init__.py Normal file
View File

@@ -0,0 +1,64 @@
"""SnapForm Job Server without Blender dependencies."""
import sys
from pathlib import Path
SERVER_DIRECTORY = Path(__file__).resolve().parent
ADDON_DIRECTORY = SERVER_DIRECTORY.parent
for import_directory in (SERVER_DIRECTORY, ADDON_DIRECTORY):
if str(import_directory) not in sys.path:
sys.path.insert(0, str(import_directory))
try:
from ..blendjob import JobServer
except ImportError:
from blendjob import JobServer
from .model_manager import ModelManager, ModelSessionCache
from .background import remove_background as _remove_background
from .da3 import generate_da3 as _generate_da3
from .media import extract_frames as _extract_frames
from .moge2 import generate_moge2 as _generate_moge2
server = JobServer("SnapForm Job Server")
@server.resource("model_manager")
def create_model_manager(job_server):
return ModelManager(job_server.storage_root / "models")
@server.job("download-model")
def download_model(context, parameters):
manager = context.resource("model_manager")
return manager.download(context, parameters["model"])
@server.job("debug-ben2")
@server.job("remove-background")
def remove_background(context, parameters):
return _remove_background(context, parameters)
@server.job("extract-frames")
def extract_frames(context, parameters):
return _extract_frames(context, parameters)
@server.job("debug-moge2")
@server.job("generate-moge2-surface")
@server.job("generate-moge2-depth")
def generate_moge2(context, parameters):
return _generate_moge2(context, parameters)
@server.job("debug-da3")
@server.job("generate-da3-surface")
@server.job("generate-da3-depth")
def generate_da3(context, parameters):
return _generate_da3(context, parameters)
__all__ = ("ModelManager", "ModelSessionCache", "server")

170
src/snapform/server/background.py Normal file
View File

@@ -0,0 +1,170 @@
import shutil
import time
from pathlib import Path
from .media import (
IMAGE_SUFFIXES,
check_cancelled,
copy_frames,
prepare_background_input,
)
def infer_foreground(
input_path,
model_dir,
device,
model_cache=None,
cancel_check=None,
):
from PIL import Image
input_path = Path(input_path)
check_cancelled(cancel_check)
if input_path.suffix.lower() not in IMAGE_SUFFIXES:
raise RuntimeError("BEN2 requires a JPG, JPEG, PNG, or WEBP image")
from .onnx_ben2 import create_session, infer
owns_model = model_cache is None
requested_device = getattr(device, "type", device)
if owns_model:
load_started = time.perf_counter()
model = create_session(model_dir, requested_device)
load_ms = (time.perf_counter() - load_started) * 1000.0
else:
model, load_ms = model_cache.get_ben2(model_dir, device)
check_cancelled(cancel_check)
with Image.open(input_path) as opened_image:
image = opened_image.convert("RGB")
inference_started = time.perf_counter()
foreground = infer(model, image)
foreground = preserve_source_rgb(image, foreground)
inference_ms = (time.perf_counter() - inference_started) * 1000.0
check_cancelled(cancel_check)
return foreground, load_ms, inference_ms
def preserve_source_rgb(source, foreground):
"""Combine the untouched source RGB with BEN2's foreground alpha."""
source_rgba = source.convert("RGBA")
alpha = foreground.convert("RGBA").getchannel("A")
if alpha.size != source_rgba.size:
from PIL import Image
alpha = alpha.resize(source_rgba.size, Image.Resampling.LANCZOS)
source_rgba.putalpha(alpha)
return source_rgba
def infer_foreground_frames(
images,
model_dir,
device,
output_directory,
progress_callback,
start_progress,
progress_span,
model_cache=None,
cancel_check=None,
):
"""Run BEN2 on every frame and write foreground.0001.png etc."""
total_load_ms = 0.0
total_inference_ms = 0.0
for index, image_path in enumerate(images):
check_cancelled(cancel_check)
progress_callback(
start_progress + progress_span * index / len(images),
f"Removing background ({index + 1}/{len(images)})",
)
foreground, load_ms, inference_ms = infer_foreground(
image_path,
model_dir,
device,
model_cache=model_cache,
cancel_check=cancel_check,
)
output_path = output_directory / f"foreground.{index + 1:04d}.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError(f"BEN2 did not write: {output_path}")
total_load_ms += load_ms
total_inference_ms += inference_ms
return total_load_ms, total_inference_ms
def remove_background(context, parameters):
"""Run BEN2 and write transparent foreground images into the Job directory."""
total_started = time.perf_counter()
context.progress(0.02, "Starting Environment")
context.check_cancelled()
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
device = parameters["device"]
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
video_frames = parameters.get("video_frames", 16)
start_frame = parameters.get("start_frame", 1)
images, temporary_directory = prepare_background_input(
parameters["input"],
video_frames,
output_directory,
start_frame=start_frame,
cancel_check=context.check_cancelled,
)
try:
context.progress(0.12, f"Loading BEN2 for {device.upper()}")
if len(images) == 1:
context.progress(0.35, "Removing background")
foreground, load_ms, inference_ms = infer_foreground(
images[0],
model_directory,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
context.progress(0.92, "Writing transparent PNG")
output_path = output_directory / "foreground.png"
foreground.convert("RGBA").save(
output_path,
format="PNG",
compress_level=1,
)
if not output_path.is_file():
raise RuntimeError("BEN2 did not write a foreground image")
else:
load_ms, inference_ms = infer_foreground_frames(
images,
model_directory,
device,
output_directory,
context.progress,
0.12,
0.7,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
copy_frames(images, output_directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
context.check_cancelled()
total_ms = (time.perf_counter() - total_started) * 1000.0
frame_text = "" if len(images) == 1 else f" from {len(images)} frame(s)"
message = (
f"Background removed{frame_text} in {inference_ms:.0f} ms "
f"(model load {load_ms:.0f} ms, total {total_ms:.0f} ms)"
)
print(message)
context.progress(1.0, message)
if tuple(output_directory.glob("foreground.*.png")):
return {"foreground": "."}
return {"foreground": "foreground.png"}

233
src/snapform/server/da3.py Normal file
View File

@@ -0,0 +1,233 @@
import shutil
from .background import infer_foreground, infer_foreground_frames
from .geometry import (
export_da3_depth_preview,
export_da3_mini_npz,
export_da3_point_cloud,
export_image_mesh_data,
export_prediction_attributes,
export_textured_mesh_glb,
)
from .media import (
copy_frames,
flatten_exported_paths,
prepare_input,
wait_for_exported_paths,
)
def merge_da3_frame_predictions(predictions):
"""Combine independent single-frame DA3 results for existing exporters."""
import numpy as np
if not predictions:
raise ValueError("DA3 returned no frame predictions")
merged = predictions[0]
for name in (
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
):
values = [getattr(prediction, name) for prediction in predictions]
setattr(merged, name, np.concatenate(values, axis=0))
return merged
def generate_da3(context, parameters):
"""Run DA3 inference and return the artifact used by the requesting Job."""
context.progress(0.02, "Starting Environment")
context.check_cancelled()
from .onnx_da3 import infer
from .onnx_runtime import provider_label
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
context.progress( 0.08, "Preparing input")
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
export_format = parameters["export_format"]
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
output_directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
context.check_cancelled()
remove_background_model = parameters.get("remove_background_model", "")
remove_background_model_dir = (
model_manager.directory(remove_background_model)
if remove_background_model
else ""
)
foreground_images = []
if remove_background_model_dir:
context.progress( 0.12, "Removing image background")
if len(images) == 1:
foreground, load_ms, inference_ms = infer_foreground(
images[0],
remove_background_model_dir,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
foreground.convert("RGBA").save(
output_directory / "foreground.png",
format="PNG",
compress_level=1,
)
if not (output_directory / "foreground.png").is_file():
raise RuntimeError("BEN2 did not write a foreground image")
foreground_images = [output_directory / "foreground.png"]
else:
load_ms, inference_ms = infer_foreground_frames(
images,
remove_background_model_dir,
device,
output_directory,
context.progress,
0.12,
0.05,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
foreground_images = sorted(
output_directory.glob("foreground.*.png")
)
if remove_background_model_dir:
print(
f"BEN2 inference {inference_ms:.0f} ms; "
f"model load {load_ms:.0f} ms"
)
context.progress( 0.15, "Loading DA3 ONNX model")
model = model_manager.get_da3(model_directory, device)
context.progress(
0.20,
f"Loaded DA3 on {provider_label(model)}",
)
context.check_cancelled()
context.progress(
0.25,
f"Estimating depth for {len(images)} frame(s)",
)
inference_options = {
"process_res": parameters["process_res"],
"process_res_method": parameters["process_res_method"],
"use_ray_pose": bool(parameters["use_ray_pose"]),
"ref_view_strategy": parameters["ref_view_strategy"],
}
try:
if parameters.get("inference_mode", "joint") == "per_frame":
predictions = []
for index, image in enumerate(images):
context.check_cancelled()
context.progress(
0.25 + 0.60 * index / len(images),
f"Estimating DA3 depth ({index + 1}/{len(images)})",
)
predictions.append(
infer(model, [str(image)], **inference_options)
)
prediction = merge_da3_frame_predictions(predictions)
else:
prediction = infer(
model,
[str(path) for path in images],
**inference_options,
)
except UnicodeDecodeError as error:
# DirectML reports failures with a locale-specific message that
# ONNX Runtime cannot always decode, masking the real error.
if provider_label(model) == "DirectML":
raise RuntimeError(
"DA3 inference failed on DirectML. Switch the device to "
"CPU in SnapForm Preferences and retry. "
f"Original error: {error}"
) from error
raise
context.check_cancelled()
context.progress( 0.88, f"Exporting {export_format}")
if export_format == "mesh":
export_textured_mesh_glb(
prediction,
images,
output_directory,
parameters.get("conf_thresh_percentile", 40.0),
bool(parameters.get("mesh_filter_edges", True)),
float(parameters.get("mesh_max_edge_length", 0.1)),
)
elif export_format == "mesh_data":
def mesh_progress(step, total, label):
context.progress(
0.88 + 0.09 * step / total,
f"Preparing surface data {step}/{total}: {label}",
)
export_image_mesh_data(
prediction,
images[0],
output_directory,
foreground_path=(
foreground_images[0] if foreground_images else None
),
progress_callback=mesh_progress,
)
elif export_format == "mini_npz":
export_da3_mini_npz(prediction, output_directory)
elif export_format == "depth_vis":
export_da3_depth_preview(prediction, output_directory)
elif export_format == "glb":
export_da3_point_cloud(
prediction,
output_directory,
parameters["conf_thresh_percentile"],
parameters["num_max_points"],
)
else:
raise ValueError(
f"DA3 ONNX does not support export format: {export_format}"
)
if export_format in {"glb", "gs_ply", "depth_vis", "mesh"}:
context.progress(
0.97,
"Writing structured debug attributes",
)
export_prediction_attributes(prediction, output_directory)
context.check_cancelled()
result_paths = wait_for_exported_paths(
output_directory,
export_format,
cancel_check=context.check_cancelled,
)
if not result_paths:
raise RuntimeError(
f"No {export_format} files were exported"
)
result_paths = flatten_exported_paths(
output_directory,
result_paths,
cancel_check=context.check_cancelled,
)
context.check_cancelled()
copy_frames(images, output_directory)
context.progress(
1.0,
f"Exported {len(result_paths)} {export_format} file(s)",
)
if context.job_type == "generate-da3-surface":
return {"surface": "mesh.npz"}
if context.job_type == "generate-da3-depth":
return {"depth": "results.npz"}
return {"directory": "."}
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

616
src/snapform/server/geometry.py Normal file
View File

@@ -0,0 +1,616 @@
from pathlib import Path
SCENE_FILENAME = "scene.glb"
def unproject_depth_frame(depth, intrinsics, extrinsics):
import numpy as np
height, width = depth.shape
x, y = np.meshgrid(
np.arange(width, dtype=np.float32),
np.arange(height, dtype=np.float32),
)
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
camera_points = rays * depth.reshape(-1, 1)
camera_points = np.column_stack(
(camera_points, np.ones(camera_points.shape[0], dtype=np.float32))
)
world_to_camera = np.eye(4, dtype=np.float32)
world_to_camera[: extrinsics.shape[0], :4] = extrinsics
world_points = (np.linalg.inv(world_to_camera) @ camera_points.T).T
return world_points[:, :3].reshape(height, width, 3)
def unproject_camera_depth(depth, intrinsics):
import numpy as np
height, width = depth.shape
x, y = np.meshgrid(
np.arange(width, dtype=np.float32),
np.arange(height, dtype=np.float32),
)
pixels = np.stack((x, y, np.ones_like(x)), axis=-1).reshape(-1, 3)
rays = (np.linalg.inv(intrinsics) @ pixels.T).T
return (rays * depth.reshape(-1, 1)).reshape(height, width, 3)
def compact_mesh(vertices, faces):
import numpy as np
vertices = np.asarray(vertices, dtype=np.float32)
faces = np.asarray(faces, dtype=np.int32)
used_vertices = np.unique(faces)
index_map = np.full(len(vertices), -1, dtype=np.int32)
index_map[used_vertices] = np.arange(len(used_vertices), dtype=np.int32)
vertices = vertices[used_vertices]
faces = index_map[faces]
return vertices, faces
def project_camera_uv(vertices, intrinsics, image_size):
import numpy as np
width, height = image_size
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Simplified mesh contains invalid camera depth")
u = intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2]
v = intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2]
denominator_x = max(width - 1, 1)
denominator_y = max(height - 1, 1)
return np.column_stack(
(
np.clip(u / denominator_x, 0.0, 1.0),
np.clip(1.0 - v / denominator_y, 0.0, 1.0),
)
).astype(np.float32)
def project_camera_pixels(vertices, intrinsics):
import numpy as np
z = vertices[:, 2]
if np.any(~np.isfinite(z)) or np.any(z <= 0.0):
raise ValueError("Mesh contains invalid camera depth")
return np.column_stack(
(
intrinsics[0, 0] * vertices[:, 0] / z + intrinsics[0, 2],
intrinsics[1, 1] * vertices[:, 1] / z + intrinsics[1, 2],
)
)
def sample_image_nearest(values, pixels):
import numpy as np
height, width = values.shape
x = np.clip(np.rint(pixels[:, 0]).astype(np.int32), 0, width - 1)
y = np.clip(np.rint(pixels[:, 1]).astype(np.int32), 0, height - 1)
return np.asarray(values)[y, x]
def confidence_percentile_map(depth, confidence):
import numpy as np
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
result = np.zeros(depth.shape, dtype=np.float32)
if not valid.any():
return result
sorted_values = np.sort(np.asarray(confidence)[valid])
ranks = np.searchsorted(
sorted_values,
np.asarray(confidence)[valid],
side="right",
)
result[valid] = ranks.astype(np.float32) / len(sorted_values)
return result
def depth_discontinuity_ratio(depth):
import numpy as np
padded = np.pad(depth, 1, mode="edge")
windows = np.lib.stride_tricks.sliding_window_view(padded, (3, 3))
difference = np.nanmax(windows, axis=(-2, -1))
difference -= np.nanmin(windows, axis=(-2, -1))
return difference / np.maximum(depth, 1e-6)
def triangle_edge_stretch_values(points, triangles):
import numpy as np
flat_points = np.asarray(points, dtype=np.float32).reshape(-1, 3)
triangle_points = flat_points[triangles]
triangle_depth = triangle_points[..., 2]
safe_depth = np.where(
np.abs(triangle_depth) > 1e-8,
triangle_depth,
np.nan,
)
triangle_rays = triangle_points / safe_depth[..., None]
next_points = np.roll(triangle_points, -1, axis=1)
next_rays = np.roll(triangle_rays, -1, axis=1)
next_depth = np.roll(triangle_depth, -1, axis=1)
local_depth = (triangle_depth + next_depth) * 0.5
flat_edges = (
triangle_rays * local_depth[..., None]
- next_rays * local_depth[..., None]
)
flat_lengths = np.linalg.norm(flat_edges, axis=-1)
actual_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
stretch = actual_lengths / np.maximum(flat_lengths, 1e-8)
stretch[np.isclose(stretch, 1.0, rtol=1e-5, atol=1e-6)] = 1.0
return stretch
def image_mesh_attributes(
vertices,
faces,
intrinsics,
confidence_percentiles,
face_depth_discontinuity,
):
import numpy as np
point_percentile = sample_image_nearest(
confidence_percentiles,
project_camera_pixels(vertices, intrinsics),
).astype(np.float32)
face_confidence = point_percentile[faces].min(axis=1)
return {
"point_depth": vertices[:, 2].astype(np.float32),
"face_confidence_percentile": face_confidence.astype(np.float32),
"face_depth_discontinuity": np.asarray(
face_depth_discontinuity,
dtype=np.float32,
),
}
def export_prediction_attributes(prediction, output_directory):
import numpy as np
values = {}
for name in ("depth", "conf", "intrinsics", "extrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.asarray(value)
if "depth" in values and "conf" in values:
values["confidence_percentile"] = np.stack(
[
confidence_percentile_map(depth, confidence)
for depth, confidence in zip(
values["depth"],
values["conf"],
)
]
)
output_path = Path(output_directory) / "prediction.attributes.npz"
np.savez(output_path, **values)
return output_path
def export_da3_mini_npz(prediction, output_directory):
import numpy as np
values = {"depth": np.round(prediction.depth, 8)}
for name in ("conf", "extrinsics", "intrinsics"):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.round(value, 2) if name == "conf" else value
output_path = Path(output_directory) / "results.npz"
np.savez_compressed(output_path, **values)
return output_path
def export_da3_depth_preview(prediction, output_directory):
import numpy as np
from PIL import Image
target = Path(output_directory) / "depth_vis"
target.mkdir(parents=True, exist_ok=True)
for index, depth in enumerate(prediction.depth):
valid = np.isfinite(depth) & (depth > 0.0)
normalized = np.zeros(depth.shape, dtype=np.float32)
if valid.any():
low, high = np.percentile(depth[valid], (2.0, 98.0))
normalized[valid] = np.clip(
(depth[valid] - low) / max(float(high - low), 1e-6),
0.0,
1.0,
)
color = np.stack(
(
normalized,
1.0 - np.abs(normalized * 2.0 - 1.0),
1.0 - normalized,
),
axis=-1,
)
color = np.asarray(color * 255.0, dtype=np.uint8)
source = np.asarray(prediction.processed_images[index], dtype=np.uint8)
combined = np.concatenate((source, color), axis=1)
Image.fromarray(combined, mode="RGB").save(
target / f"{index:04d}.jpg",
quality=95,
)
def export_da3_point_cloud(
prediction,
output_directory,
confidence_percentile=40.0,
num_max_points=1_000_000,
):
import numpy as np
import trimesh
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
threshold = confidence_percentile_threshold(
depth,
confidence,
confidence_percentile,
)
point_sets = []
color_sets = []
for index in range(depth.shape[0]):
points = unproject_depth_frame(
depth[index],
prediction.intrinsics[index],
prediction.extrinsics[index],
)
valid = np.isfinite(points).all(axis=-1) & (depth[index] > 0.0)
if threshold is not None:
valid &= confidence[index] >= threshold
point_sets.append(points[valid])
color_sets.append(prediction.processed_images[index][valid])
points = np.concatenate(point_sets, axis=0)
colors = np.concatenate(color_sets, axis=0)
maximum = max(int(num_max_points), 1)
if len(points) > maximum:
indices = np.linspace(0, len(points) - 1, maximum, dtype=np.int64)
points = points[indices]
colors = colors[indices]
scene = trimesh.Scene(trimesh.points.PointCloud(points, colors=colors))
scene.export(Path(output_directory) / SCENE_FILENAME)
def remove_background_faces(faces, foreground_mask):
"""Return only triangles whose vertices all belong to the foreground."""
foreground_points = foreground_mask.reshape(-1)
return faces[foreground_points[faces].all(axis=1)]
def export_image_mesh_data(
prediction,
image_path,
output_directory,
foreground_path=None,
progress_callback=None,
):
import numpy as np
from PIL import Image
total_steps = 6
def progress(step, label):
if progress_callback is not None:
progress_callback(step, total_steps, label)
progress(1, "Validating DA3 data")
required = ("depth", "conf", "intrinsics")
missing = [
name
for name in required
if getattr(prediction, name, None) is None
]
if missing:
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
depth_frames = np.asarray(prediction.depth, dtype=np.float32)
confidence_frames = np.asarray(prediction.conf, dtype=np.float32)
if depth_frames.shape[0] != 1:
raise ValueError("Convert to Surface currently supports one image")
raw_depth = depth_frames[0]
confidence = confidence_frames[0]
intrinsics = np.asarray(prediction.intrinsics[0], dtype=np.float32)
progress(2, "Unprojecting the depth grid")
points = unproject_camera_depth(raw_depth, intrinsics)
progress(3, "Preparing confidence and background data")
foreground_mask = np.ones(raw_depth.shape, dtype=bool)
if foreground_path is not None:
foreground_mask = foreground_mesh_mask(
foreground_path,
raw_depth.shape,
)
confidence_percentiles = confidence_percentile_map(
raw_depth,
confidence,
)
progress(4, "Building source triangles")
faces = mesh_triangle_faces(
points,
raw_depth,
confidence,
None,
False,
0.0,
)
if foreground_path is not None:
faces = remove_background_faces(faces, foreground_mask)
if not len(faces):
raise ValueError(
"DA3 output does not contain any valid surface faces"
)
progress(5, "Preparing source UV and attributes")
depth_discontinuity = depth_discontinuity_ratio(raw_depth)
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
vertices, faces = compact_mesh(
points.reshape(-1, 3),
faces,
)
height, width = raw_depth.shape
uv = project_camera_uv(vertices, intrinsics, (width, height))
attributes = image_mesh_attributes(
vertices,
faces,
intrinsics,
confidence_percentiles,
face_depth_discontinuity,
)
reference_depth = float(np.median(vertices[:, 2]))
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
progress(6, "Writing source surface data")
result_path = output_directory / "mesh.npz"
np.savez_compressed(
result_path,
vertices=vertices,
faces=faces,
uv=uv,
intrinsics=intrinsics,
image_size=np.asarray((width, height), dtype=np.int32),
reference_depth=np.asarray(reference_depth, dtype=np.float32),
**attributes,
)
texture_path = output_directory / "texture.png"
texture_source = foreground_path or image_path
with Image.open(texture_source) as image:
image.convert("RGBA").save(texture_path, format="PNG", compress_level=1)
return result_path, texture_path
def depth_edge_mask(depth, threshold=12.0 / 255.0):
import numpy as np
values = np.asarray(depth, dtype=np.float32)
edges = np.zeros(values.shape, dtype=bool)
for axis in (0, 1):
first = np.take(values, range(values.shape[axis] - 1), axis=axis)
second = np.take(values, range(1, values.shape[axis]), axis=axis)
valid = (
np.isfinite(first)
& np.isfinite(second)
& (first > 0.0)
& (second > 0.0)
)
scale = np.maximum(np.minimum(np.abs(first), np.abs(second)), 1e-6)
discontinuity = valid & (np.abs(first - second) / scale > threshold)
lower = [slice(None), slice(None)]
upper = [slice(None), slice(None)]
lower[axis] = slice(0, -1)
upper[axis] = slice(1, None)
edges[tuple(lower)] |= discontinuity
edges[tuple(upper)] |= discontinuity
return edges
def edge_culling_stretch_limit(culling):
factor = min(max(float(culling) / 100.0, 0.0001), 1.0)
return 1.0 / factor
def foreground_mesh_mask(foreground_path, depth_shape):
import numpy as np
from PIL import Image, ImageFilter
height, width = depth_shape
with Image.open(foreground_path) as source:
alpha = source.convert("RGBA").getchannel("A")
alpha = alpha.resize(
(width, height),
resample=Image.Resampling.BILINEAR,
)
alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.75))
return np.asarray(alpha) >= 128
def confidence_percentile_threshold(depth, confidence, percentile):
import numpy as np
if float(percentile) <= 0.0:
return None
valid = np.isfinite(depth) & (depth > 0.0) & np.isfinite(confidence)
if not valid.any():
raise ValueError("DA3 returned no valid depth confidence values")
return float(np.percentile(confidence[valid], percentile))
def mesh_triangle_faces(
points,
depth,
confidence,
confidence_threshold,
filter_edges,
max_edge_length,
edge_threshold=12.0 / 255.0,
edge_stretch_limit=None,
vertex_mask=None,
):
import numpy as np
height, width = depth.shape
indices = np.arange(height * width, dtype=np.int32).reshape(height, width)
quads = np.stack(
(
indices[:-1, :-1],
indices[1:, :-1],
indices[1:, 1:],
indices[:-1, 1:],
),
axis=-1,
).reshape(-1, 4)
valid = (
np.isfinite(points).all(axis=-1)
& np.isfinite(depth)
& (depth > 0.0)
)
if confidence_threshold is not None:
valid &= np.isfinite(confidence) & (
confidence >= confidence_threshold
)
if filter_edges:
valid &= ~depth_edge_mask(depth, edge_threshold)
if vertex_mask is not None:
if np.shape(vertex_mask) != depth.shape:
raise ValueError("Foreground mask and depth dimensions do not match")
valid &= np.asarray(vertex_mask, dtype=bool)
triangles = np.concatenate(
(quads[:, (0, 1, 2)], quads[:, (0, 2, 3)]),
axis=0,
)
face_valid = valid.reshape(-1)[triangles].all(axis=1)
if edge_stretch_limit is not None:
face_valid &= triangle_edge_stretch_mask(
points,
triangles,
edge_stretch_limit,
)
if max_edge_length > 0.0:
flat_points = points.reshape(-1, 3)
triangle_points = flat_points[triangles]
next_points = np.roll(triangle_points, -1, axis=1)
edge_lengths = np.linalg.norm(
triangle_points - next_points,
axis=-1,
)
face_valid &= (edge_lengths <= max_edge_length).all(axis=1)
triangles = triangles[face_valid]
if not len(triangles):
return np.empty((0, 3), dtype=np.int32)
return triangles
def triangle_edge_stretch_mask(points, triangles, stretch_limit):
stretch = triangle_edge_stretch_values(points, triangles)
tolerance = float(stretch_limit) * (1.0 + 1e-5)
return (stretch <= tolerance).all(axis=1)
def export_textured_mesh_glb(
prediction,
image_paths,
output_directory,
confidence_percentile,
filter_edges,
max_edge_length,
):
import numpy as np
import trimesh
from PIL import Image
required = ("depth", "conf", "intrinsics", "extrinsics")
missing = [name for name in required if getattr(prediction, name) is None]
if missing:
raise ValueError(f"DA3 mesh output is missing: {', '.join(missing)}")
depth = np.asarray(prediction.depth, dtype=np.float32)
confidence = np.asarray(prediction.conf, dtype=np.float32)
threshold = confidence_percentile_threshold(
depth,
confidence,
confidence_percentile,
)
meshes = []
referenced_points = []
for index in range(depth.shape[0]):
points = unproject_depth_frame(
depth[index],
prediction.intrinsics[index],
prediction.extrinsics[index],
)
faces = mesh_triangle_faces(
points,
depth[index],
confidence[index],
threshold,
filter_edges,
max_edge_length,
)
if not len(faces):
continue
height, width = depth[index].shape
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uu, vv = np.meshgrid(u, v)
uv = np.stack((uu, vv), axis=-1).reshape(-1, 2)
with Image.open(image_paths[index]) as source_image:
texture_image = source_image.convert("RGB").copy()
visual = trimesh.visual.texture.TextureVisuals(
uv=uv,
image=texture_image,
)
mesh = trimesh.Trimesh(
vertices=points.reshape(-1, 3),
faces=faces,
visual=visual,
process=False,
)
mesh.vertex_attributes["o_depth"] = depth[index].reshape(-1)
mesh.vertex_attributes["o_confidence"] = confidence[
index
].reshape(-1)
mesh.vertex_attributes["o_confidence_percentile"] = (
confidence_percentile_map(
depth[index],
confidence[index],
).reshape(-1)
)
mesh.remove_unreferenced_vertices()
meshes.append(mesh)
referenced_points.append(mesh.vertices)
if not meshes:
raise ValueError(
"Surface filters removed every face; lower Confidence Threshold or disable filters"
)
world_to_camera = np.asarray(prediction.extrinsics[0], dtype=np.float64)
if world_to_camera.shape == (3, 4):
homogeneous = np.eye(4, dtype=np.float64)
homogeneous[:3, :4] = world_to_camera
world_to_camera = homogeneous
cv_to_gltf = np.eye(4, dtype=np.float64)
cv_to_gltf[1, 1] = -1.0
cv_to_gltf[2, 2] = -1.0
alignment = cv_to_gltf @ world_to_camera
aligned_points = trimesh.transform_points(
np.concatenate(referenced_points, axis=0),
alignment,
)
center = np.median(aligned_points, axis=0)
alignment[:3, 3] -= center
scene = trimesh.Scene()
for index, mesh in enumerate(meshes):
mesh.apply_transform(alignment)
scene.add_geometry(mesh, node_name=f"DepthMesh.{index + 1:04d}")
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
output_path = output_directory / "mesh.glb"
scene.export(output_path)
return output_path

306
src/snapform/server/media.py Normal file
View File

@@ -0,0 +1,306 @@
import shutil
import tempfile
import time
from pathlib import Path
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".webp"}
VIDEO_SUFFIXES = {
".mp4",
".mov",
".avi",
".mkv",
".webm",
".m4v",
".mts",
".mpg",
".mpeg",
}
SCENE_FILENAME = "scene.glb"
class TaskCancelled(RuntimeError):
pass
def check_cancelled(cancel_check):
if cancel_check is not None and cancel_check():
raise TaskCancelled("Task cancelled")
def collect_images(directory):
directory = Path(directory)
return sorted(
path
for path in directory.iterdir()
if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
)
def first_frames(images, maximum_frames, start_frame=1):
"""Return up to maximum_frames image paths in order from start_frame."""
if maximum_frames <= 0:
return images
start_index = max(start_frame - 1, 0)
return images[start_index:start_index + maximum_frames]
def extract_video_frames(
video_path,
output_directory,
maximum_frames,
av,
start_frame=1,
cancel_check=None,
):
output_directory = Path(output_directory)
output_directory.mkdir(parents=True, exist_ok=True)
try:
container = av.open(str(video_path))
except Exception as error:
raise RuntimeError(f"Unable to open video: {video_path}") from error
image_paths = []
skipped = 0
try:
stream = container.streams.video[0]
for frame in container.decode(stream):
check_cancelled(cancel_check)
if skipped < start_frame - 1:
skipped += 1
continue
image_path = output_directory / f"frame.{len(image_paths) + 1:04d}.jpg"
frame.to_image().save(image_path, format="JPEG")
image_paths.append(image_path)
if len(image_paths) >= maximum_frames:
break
finally:
container.close()
if not image_paths:
raise RuntimeError(f"Unable to read any video frame from: {video_path}")
return image_paths
def prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
check_cancelled(cancel_check)
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
suffix = input_path.suffix.lower()
if suffix in IMAGE_SUFFIXES:
return [input_path], None
if suffix not in VIDEO_SUFFIXES:
supported = ", ".join(sorted(IMAGE_SUFFIXES | VIDEO_SUFFIXES))
raise RuntimeError(f"Unsupported input type; expected one of: {supported}")
import av
temporary_directory = Path(
tempfile.mkdtemp(prefix="snapform-video-", dir=temporary_root)
)
try:
images = extract_video_frames(
input_path,
temporary_directory,
video_frames,
av,
start_frame=start_frame,
cancel_check=cancel_check,
)
except Exception:
shutil.rmtree(temporary_directory, ignore_errors=True)
raise
return images, temporary_directory
def prepare_background_input(
input_path,
video_frames,
temporary_root,
start_frame=1,
cancel_check=None,
):
"""Resolve a BEN2 input to frame images; returns (images, temporary_directory)."""
input_path = Path(input_path)
if input_path.is_dir():
images = collect_images(input_path)
if not images:
raise RuntimeError(f"No image frames found in: {input_path}")
return first_frames(images, video_frames, start_frame), None
if input_path.suffix.lower() in IMAGE_SUFFIXES:
return [input_path], None
return prepare_input(
input_path,
video_frames,
temporary_root,
start_frame=start_frame,
cancel_check=cancel_check,
)
def copy_frames(images, output_directory):
"""Copy a multi-frame input as numbered frame files for the Blender sequence."""
if len(images) <= 1:
return
for index, image_path in enumerate(images):
suffix = Path(image_path).suffix.lower() or ".jpg"
destination = (
Path(output_directory) / f"frame.{index + 1:04d}{suffix}"
)
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def extract_frames(context, parameters):
"""Copy selected source frames into the current job directory."""
context.directory.mkdir(parents=True, exist_ok=True)
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
context.directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
copy_frames(images, context.directory)
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)
message = f"Extracted {len(images)} frame(s)"
context.progress(1.0, message)
return {"frames": len(images)}
def exported_paths(output_directory, export_format):
output_directory = Path(output_directory)
if export_format == "glb":
scene_path = output_directory / SCENE_FILENAME
return [scene_path] if scene_path.is_file() else []
if export_format == "gs_ply":
return sorted((output_directory / "gs_ply").glob("*.ply"))
if export_format == "depth_vis":
return sorted((output_directory / "depth_vis").glob("*.jpg"))
if export_format == "mesh":
result_path = output_directory / "mesh.glb"
return [result_path] if result_path.is_file() else []
if export_format == "mesh_data":
result_path = output_directory / "mesh.npz"
return [result_path] if result_path.is_file() else []
if export_format == "mini_npz":
result_path = output_directory / "results.npz"
if result_path.is_file():
return [result_path]
result_path = (
output_directory / "exports" / "mini_npz" / "results.npz"
)
return [result_path] if result_path.is_file() else []
return []
def npz_is_complete(path):
"""Return True only when the npz archive can be fully read."""
try:
import numpy as np
with np.load(path, allow_pickle=False) as archive:
return "depth" in archive.files
except Exception:
return False
def wait_for_exported_paths(
output_directory,
export_format,
timeout=30.0,
cancel_check=None,
):
# DA3 的 NPZ exporter 在后台线程写入，文件出现时可能尚未写完整，
# 直接返回会让下游读到半截 zip 而报 EOFError，因此需等文件可完整读取
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
check_cancelled(cancel_check)
result_paths = exported_paths(output_directory, export_format)
if result_paths:
if (
export_format in ("mini_npz", "npz")
and not npz_is_complete(result_paths[0])
):
time.sleep(0.1)
continue
return result_paths
time.sleep(0.05)
return []
def move_exported_file(
source,
destination,
timeout=10.0,
cancel_check=None,
):
"""Move an async export after Windows releases its file handle."""
source = Path(source)
destination = Path(destination)
deadline = time.monotonic() + timeout
while True:
check_cancelled(cancel_check)
try:
source.replace(destination)
return destination
except OSError as error:
is_file_lock = isinstance(error, PermissionError) or getattr(
error,
"winerror",
None,
) in {32, 33}
if not is_file_lock or time.monotonic() >= deadline:
raise
time.sleep(0.1)
def flatten_exported_paths(
output_directory,
result_paths,
cancel_check=None,
):
"""Move files produced in exporter subdirectories to the job root."""
output_directory = Path(output_directory)
flattened_paths = []
for source in result_paths:
source = Path(source)
check_cancelled(cancel_check)
if source.parent == output_directory:
flattened_paths.append(source)
continue
destination = output_directory / source.name
if destination.exists():
raise RuntimeError(f"Export filename already exists: {destination.name}")
flattened_paths.append(
move_exported_file(
source,
destination,
cancel_check=cancel_check,
)
)
for directory in sorted(
{path.parent for path in map(Path, result_paths)},
key=lambda path: len(path.parts),
reverse=True,
):
while directory != output_directory:
try:
directory.rmdir()
except OSError:
break
directory = directory.parent
return flattened_paths

136
src/snapform/models.py → src/snapform/server/model_catalog.py
View File

@@ -153,125 +153,6 @@ MODEL_NUMBERS = {
}
_ENUM_ITEMS_CACHE = {}
_DEPTH_PLANE_ENUM_ITEMS_CACHE = {}
_SURFACE_ENUM_ITEMS_CACHE = {}
def enum_items(supports_gs=None, show_download_status=False):
model_ready = None
if show_download_status:
from .client import model_ready
models = [
model
for model in MODELS.values()
if model.enabled
and (supports_gs is None or model.supports_gs == supports_gs)
]
download_status = (
tuple(model_ready(model.key) for model in models)
if model_ready is not None
else ()
)
cache_key = (
supports_gs,
show_download_status,
download_status,
)
cached_items = _ENUM_ITEMS_CACHE.get(cache_key)
if cached_items is not None:
return cached_items
items = []
for index, model in enumerate(models):
label = model.label
if download_status and not download_status[index]:
label = f"{label} (Missing)"
items.append(
(
model.key,
label,
f"{model.description}; license: {model.license_name}",
MODEL_NUMBERS[model.key],
)
)
_ENUM_ITEMS_CACHE[cache_key] = items
return items
def model_enum_items(_owner, _context):
return enum_items(show_download_status=True)
def gs_model_enum_items(_owner, _context):
return enum_items(supports_gs=True, show_download_status=True)
def moge2_model_enum_items(_owner, _context):
from .client import model_ready
return [
(
model.key,
model.label if model_ready(model.key) else f"{model.label} (Missing)",
f"{model.description}; license: {model.license_name}",
index,
)
for index, model in enumerate(MOGE2_MODELS.values())
]
def depth_plane_model_enum_items(_owner, _context):
from .client import model_ready
models = [
*MOGE2_MODELS.values(),
*(model for model in MODELS.values() if model.enabled),
]
ready = tuple(model_ready(model.key) for model in models)
cached = _DEPTH_PLANE_ENUM_ITEMS_CACHE.get(ready)
if cached is not None:
return cached
items = []
for index, model in enumerate(models):
label = model.label if ready[index] else f"{model.label} (Missing)"
items.append(
(
model.key,
label,
f"{model.description}; license: {model.license_name}",
index,
)
)
items.append(("ORIGINAL", "Original", "Use the original image as depth", len(items)))
_DEPTH_PLANE_ENUM_ITEMS_CACHE[ready] = items
return items
def surface_model_enum_items(owner, context):
depth_items = depth_plane_model_enum_items(owner, context)
cache_key = id(depth_items)
cached = _SURFACE_ENUM_ITEMS_CACHE.get(cache_key)
if cached is not None:
return cached
items = [
item
for item in depth_items
if item[0] != "ORIGINAL"
]
_SURFACE_ENUM_ITEMS_CACHE[cache_key] = items
return items
def selected_model_key(settings, export_format=None):
selected_format = export_format or settings.export_format
if selected_format == "gs_ply":
return settings.gs_model
return settings.model
def get_model(model_key):
try:
return MODELS[model_key]
@@ -289,3 +170,20 @@ def get_downloadable_model(model_key):
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
def model_record(model, ready=False):
"""Return public model metadata exposed by the Models Resource."""
return {
"key": model.key,
"label": model.label,
"description": model.description,
"license": model.license_name,
"supports_gs": model.supports_gs,
"enabled": model.enabled,
"family": "moge2" if model.key.startswith("MOGE2_") else (
"ben2" if model.key == BEN2_MODEL.key else "da3"
),
"ready": bool(ready),
}

444
src/snapform/server/model_download.py Normal file
View File

@@ -0,0 +1,444 @@
import hashlib
import http.client
import os
import threading
import time
from pathlib import Path
from urllib.parse import quote
from urllib.request import (
HTTPSHandler,
ProxyHandler,
Request,
build_opener,
urlopen,
)
HF_ENDPOINT = "https://huggingface.co"
R2_ENDPOINT = "https://models.omoolab.xyz"
CONNECT_TIMEOUT = 3.0
DOWNLOAD_CHUNK_SIZE = 256 * 1024
class _ShortConnectHTTPSConnection(http.client.HTTPSConnection):
"""Use a short timeout for connection setup, without limiting downloads."""
def connect(self):
try:
self.timeout = CONNECT_TIMEOUT
super().connect()
finally:
self.timeout = None
self.sock.settimeout(None)
class _ShortConnectHTTPSHandler(HTTPSHandler):
def https_open(self, request):
return self.do_open(_ShortConnectHTTPSConnection, request)
_download_opener = build_opener(ProxyHandler(), _ShortConnectHTTPSHandler())
class DownloadProgressReporter:
def __init__(
self,
progress_callback,
repository,
clock=time.monotonic,
start=0.05,
span=0.90,
label=None,
):
self.progress_callback = progress_callback
self.repository = repository
self.clock = clock
self.lock = threading.Lock()
self.start = start
self.span = span
self.label = label
self.progress = start
self.last_write_time = 0.0
self.started_at = self.clock()
def report(self, current, total, rate=None):
if not total or current <= 0:
return
factor = min(max(current / total, 0.0), 1.0)
progress = self.start + factor * self.span
now = self.clock()
with self.lock:
if progress <= self.progress:
return
if now - self.last_write_time < 0.20 and factor < 1.0:
return
self.progress = progress
self.last_write_time = now
effective_rate = rate
elapsed = now - self.started_at
if not effective_rate and elapsed > 0:
effective_rate = current / elapsed
speed = format_rate(effective_rate) if effective_rate else ""
prefix = self.label or self.repository
message = f"Downloading {prefix}"
if speed:
message = f"{message} {speed}"
self.progress_callback(
progress,
message,
)
def create_download_progress_class(
progress_callback,
repository,
tqdm_type=None,
clock=time.monotonic,
start=0.05,
span=0.90,
label=None,
):
if tqdm_type is None:
from tqdm.auto import tqdm as tqdm_type
reporter = DownloadProgressReporter(
progress_callback,
repository,
clock,
start,
span,
label,
)
class DownloadProgress(tqdm_type):
def __init__(self, *args, **kwargs):
description = str(kwargs.get("desc", ""))
# huggingface_hub 1.x 的字节进度条描述是 "Downloading (incomplete total...)"，
# 用前缀匹配而非精确匹配，否则真实下载进度会被忽略。
self.snapform_reports_progress = (
description.startswith("Downloading")
or description.startswith("Fetching ")
)
super().__init__(*args, **kwargs)
def update(self, amount=1):
result = super().update(amount)
if self.snapform_reports_progress:
reporter.report(self.n, self.total, self._download_rate())
return result
def refresh(self, *args, **kwargs):
result = super().refresh(*args, **kwargs)
if self.snapform_reports_progress:
reporter.report(self.n, self.total, self._download_rate())
return result
def _download_rate(self):
try:
return getattr(self, "format_dict", {}).get("rate")
except Exception:
return None
return DownloadProgress
def progress_range(args):
start = getattr(args, "progress_start", None)
end = getattr(args, "progress_end", None)
if start is not None and end is not None:
return float(start), float(end) - float(start)
return 0.05, 0.90
def download_model(args):
label = getattr(args, "label", None) or args.repository
model_dir = Path(args.model_dir)
model_dir.mkdir(parents=True, exist_ok=True)
start, span = progress_range(args)
r2_directory = getattr(args, "r2_directory", "")
r2_files = getattr(args, "r2_file", [])
if r2_directory and r2_files:
_download_declared_or_fail(
args,
model_dir,
label,
start,
span,
)
return
_download_snapshot_or_fail(args, model_dir, label, start, span)
def _download_declared_or_fail(args, model_dir, label, start, span):
"""Download declared files from Hugging Face, then the R2 mirror."""
sources = (
(HF_ENDPOINT, download_from_hf, f"Downloading {label}"),
(
R2_ENDPOINT,
download_from_r2,
f"Downloading {label} (mirror)",
),
)
failures = []
for source, downloader, message in sources:
args.progress(start, message)
try:
downloader(
args,
model_dir,
start=start,
span=span,
label=label,
)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
cancel_check = getattr(args, "cancel_check", None)
if cancel_check is not None:
cancel_check()
failures.append((source, error))
continue
args.progress(start + span, f"{label} ready")
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
raise RuntimeError(f"All download sources failed: {detail}")
def _download_snapshot_or_fail(args, model_dir, label, start, span):
"""Fallback download for repositories without declared files, via huggingface_hub."""
failures = []
try:
from huggingface_hub import snapshot_download
except Exception as error:
snapshot_download = None
failures.append((HF_ENDPOINT, error))
progress_class = None
if snapshot_download:
progress_class = create_download_progress_class(
args.progress,
args.repository,
start=start,
span=span,
label=label,
)
for endpoint in download_endpoints() if snapshot_download else ():
os.environ["HF_ENDPOINT"] = endpoint
# snapshot_download 先拉取文件清单并校验 ETag，还没有字节下载，
# 用 Checking 提示当前阶段，避免进度条看起来卡住。
message = f"Checking {label} files"
if endpoint != HF_ENDPOINT:
message = f"{message} (mirror)"
args.progress(start, message)
try:
snapshot_download(
repo_id=args.repository,
local_dir=model_dir,
tqdm_class=progress_class,
allow_patterns=args.allow_pattern or None,
endpoint=endpoint,
)
_verify_model_complete(
model_dir,
getattr(args, "required_pattern", None),
)
except Exception as error:
cancel_check = getattr(args, "cancel_check", None)
if cancel_check is not None:
cancel_check()
failures.append((endpoint, error))
continue
args.progress(start + span, f"{label} ready")
return
detail = "; ".join(
f"{source}: {error}" for source, error in failures
)
raise RuntimeError(f"All download sources failed: {detail}")
def download_endpoints():
"""Return the configured Hugging Face endpoint or official source."""
endpoint = os.environ.get("HF_ENDPOINT") or HF_ENDPOINT
return [endpoint.rstrip("/")]
def _verify_model_complete(model_dir, required_patterns=None):
"""Refuse to report success when the required model files are missing."""
model_dir = Path(model_dir)
required_patterns = required_patterns or (
"config.json",
"*.safetensors",
)
missing = [
pattern
for pattern in required_patterns
if not any(model_dir.glob(pattern))
]
if missing:
raise RuntimeError(
"Model download reported success but files are missing: "
+ ", ".join(missing)
)
def _download_file_from_url(
url,
destination,
expected_size,
expected_checksum,
progress=None,
retries=1,
cancel_check=None,
):
"""Download and verify one model file atomically from a direct URL."""
request = Request(url, headers={"User-Agent": "SnapForm"})
opener = _download_opener.open if request.type == "https" else urlopen
last_error = None
for _attempt in range(retries):
temporary = destination.with_name(f"{destination.name}.part")
try:
checksum = hashlib.sha256()
downloaded = 0
with opener(request) as response:
with open(temporary, "wb") as output:
while True:
if cancel_check is not None:
cancel_check()
chunk = response.read(DOWNLOAD_CHUNK_SIZE)
if not chunk:
break
output.write(chunk)
checksum.update(chunk)
downloaded += len(chunk)
if progress is not None:
progress(len(chunk))
if downloaded != expected_size:
raise RuntimeError(
f"Size mismatch for {destination.name}: "
f"{downloaded} != {expected_size}"
)
if checksum.hexdigest().lower() != expected_checksum.lower():
raise RuntimeError(f"SHA-256 mismatch for {destination.name}")
temporary.replace(destination)
return
except Exception as error:
if cancel_check is not None:
cancel_check()
last_error = error
try:
temporary.unlink()
except OSError:
pass
raise last_error
def r2_file_url(directory, filename):
return f"{R2_ENDPOINT}/{quote(directory)}/{quote(filename)}"
def hf_file_url(repository, filename):
"""Hugging Face resolve URL for a direct, metadata-free download."""
endpoint = (os.environ.get("HF_ENDPOINT") or HF_ENDPOINT).rstrip("/")
return f"{endpoint}/{repository}/resolve/main/{quote(filename)}"
def _parse_declared_files(args):
selected = []
for value in args.r2_file:
filename, size, checksum = value.split("|", 2)
selected.append((filename, int(size), checksum))
return selected
def _download_declared_files(
args,
model_dir,
url_builder,
start,
span,
label,
retries=1,
):
"""Download each declared file through url_builder with aggregated progress."""
repository = args.repository
selected = _parse_declared_files(args)
reporter = DownloadProgressReporter(
args.progress,
repository,
start=start,
span=span,
label=label or repository,
)
total_size = sum(size for _path, size, _checksum in selected)
downloaded = 0
def on_chunk(count):
nonlocal downloaded
downloaded += count
reporter.report(downloaded, total_size)
for path, size, checksum in selected:
destination = Path(model_dir) / path
destination.parent.mkdir(parents=True, exist_ok=True)
_download_file_from_url(
url_builder(path),
destination,
size,
checksum,
progress=on_chunk,
retries=retries,
cancel_check=getattr(args, "cancel_check", None),
)
def download_from_r2(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
):
"""Download the declared model files from OmooLab R2."""
_download_declared_files(
args,
model_dir,
lambda filename: r2_file_url(args.r2_directory, filename),
start,
span,
f"{label} (mirror)" if label else "mirror",
)
def download_from_hf(
args,
model_dir,
start=0.05,
span=0.90,
label=None,
):
"""Download the declared model files directly from Hugging Face resolve URLs."""
_download_declared_files(
args,
model_dir,
lambda filename: hf_file_url(args.repository, filename),
start,
span,
label,
)
def format_rate(rate):
size = float(rate)
if size <= 0:
return ""
units = ("B/s", "KB/s", "MB/s", "GB/s")
unit_index = 0
while size >= 1024 and unit_index < len(units) - 1:
size /= 1024
unit_index += 1
if unit_index == 0:
return f"{size:.0f} {units[unit_index]}"
return f"{size:.1f} {units[unit_index]}"

50
src/snapform/ai_service/model_manager.py → src/snapform/server/model_manager.py
View File

@@ -3,12 +3,12 @@ import time
from pathlib import Path
from types import SimpleNamespace
import jobs
import models as model_catalog
from . import model_catalog
from .model_download import download_model
class ModelSessionCache:
"""Cache ONNX sessions shared by jobs in one service process."""
"""Cache ONNX sessions shared by Jobs in one Server process."""
def __init__(self):
self.ben2_model = None
@@ -27,7 +27,7 @@ class ModelSessionCache:
return self._load_ben2(model_dir, requested_device, key)
def _load_ben2(self, model_dir, requested_device, key):
from ben2_onnx import create_session
from .onnx_ben2 import create_session
started = time.perf_counter()
model = create_session(model_dir, requested_device)
@@ -47,15 +47,15 @@ class ModelSessionCache:
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, requested_device, key):
from da3_onnx import create_session
from .onnx_da3 import create_session
model = create_session(model_dir, requested_device)
self.da3_model = model
self.da3_key = key
return model
def get_moge_onnx(self, model_dir, requested_device):
from moge_onnx import create_session
def get_moge2(self, model_dir, requested_device):
from .onnx_moge2 import create_session
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None:
@@ -112,7 +112,27 @@ class ModelManager:
def directory(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
return self.models_directory / model.directory_name
directory = self.models_directory / model.directory_name
old_moge_directory = self.models_directory / "moge2_base_normal"
if (
model_key == "MOGE2_VITB_NORMAL"
and not directory.is_dir()
and old_moge_directory.is_dir()
):
old_moge_directory.rename(directory)
legacy_directory = self.models_directory / model_key.lower()
if (
not directory.is_dir()
and legacy_directory.is_dir()
and legacy_directory != directory
):
legacy_directory.rename(directory)
return directory
def ready(self, model_key):
model = model_catalog.get_downloadable_model(model_key)
directory = self.directory(model_key)
return all(any(directory.glob(pattern)) for pattern in model.ready_patterns)
def get_ben2(self, model_dir, device):
return self.sessions.get_ben2(model_dir, device)
@@ -120,11 +140,15 @@ class ModelManager:
def get_da3(self, model_dir, device):
return self.sessions.get_da3(model_dir, device)
def get_moge_onnx(self, model_dir, device):
return self.sessions.get_moge_onnx(model_dir, device)
def get_moge2(self, model_dir, device):
return self.sessions.get_moge2(model_dir, device)
def snapshot(self):
return {"loaded": self.sessions.snapshot()}
catalog = [
model_catalog.model_record(model, self.ready(model.key))
for model in model_catalog.DOWNLOADABLE_MODELS.values()
]
return {"loaded": self.sessions.snapshot(), "catalog": catalog}
def clear(self):
self.sessions.close()
@@ -137,7 +161,7 @@ class ModelManager:
args = SimpleNamespace(
repository=model.repository,
model_dir=str(self.directory(model_key)),
status=context.update,
progress=context.progress,
cancel_check=context.check_cancelled,
label=model.label,
allow_pattern=list(model.download_patterns),
@@ -149,5 +173,5 @@ class ModelManager:
],
)
context.check_cancelled()
jobs.download_model(args)
download_model(args)
context.check_cancelled()

304
src/snapform/server/moge2.py Normal file
View File

@@ -0,0 +1,304 @@
import shutil
from pathlib import Path
from .background import infer_foreground, infer_foreground_frames
from .geometry import (
depth_discontinuity_ratio,
depth_edge_mask,
foreground_mesh_mask,
mesh_triangle_faces,
)
from .media import copy_frames, prepare_input
def _save_moge_preview(output_directory, index, depth, normal, mask):
import numpy as np
from PIL import Image
suffix = "" if index is None else f".{index + 1:04d}"
valid = mask & np.isfinite(depth) & (depth > 0)
preview = np.zeros(depth.shape, dtype=np.uint8)
if np.any(valid):
low, high = np.percentile(depth[valid], (2.0, 98.0))
span = max(float(high - low), 1e-6)
preview[valid] = np.clip(
(depth[valid] - low) / span * 255.0,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(preview, mode="L").save(
output_directory / f"depth{suffix}.png"
)
normal_rgb = np.clip(
(normal * np.array([1.0, -1.0, -1.0], dtype=np.float32) + 1.0)
* 127.5,
0.0,
255.0,
).astype(np.uint8)
Image.fromarray(normal_rgb, mode="RGB").save(
output_directory / f"normal{suffix}.png"
)
Image.fromarray(mask.astype(np.uint8) * 255, mode="L").save(
output_directory / f"mask{suffix}.png"
)
def export_moge_debug_mesh(
output_directory,
points,
depth,
mask,
normal,
image,
export_format,
):
import numpy as np
import trimesh
clean_mask = mask & ~depth_edge_mask(depth, threshold=0.04)
faces = mesh_triangle_faces(
points,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=clean_mask,
)
vertices = points.reshape(-1, 3)
colors = np.asarray(image.reshape(-1, 3) * 255.0, dtype=np.uint8)
vertex_normals = normal.reshape(-1, 3)
vertices = vertices * np.array([1.0, -1.0, -1.0], dtype=np.float32)
vertex_normals = vertex_normals * np.array([1.0, -1.0, -1.0], dtype=np.float32)
mesh = trimesh.Trimesh(
vertices=vertices,
faces=faces,
vertex_colors=colors,
vertex_normals=vertex_normals,
process=False,
)
mesh.remove_unreferenced_vertices()
filename = "mesh.glb" if export_format == "glb" else "pointcloud.ply"
mesh.export(Path(output_directory) / filename)
def export_moge_surface_archive(
output_path,
points,
depth,
mask,
normal,
intrinsics,
foreground_mask=None,
):
import numpy as np
height, width = depth.shape
finite_mask = np.isfinite(depth) & (depth > 0.0)
finite_mask &= np.isfinite(points).all(axis=-1)
if foreground_mask is not None:
finite_mask &= foreground_mask
depth_discontinuity = depth_discontinuity_ratio(depth)
faces = mesh_triangle_faces(
points,
depth,
np.ones(depth.shape, dtype=np.float32),
None,
False,
0.0,
vertex_mask=finite_mask,
)
vertices = points.reshape(-1, 3)
u = np.linspace(0.0, 1.0, width, dtype=np.float32)
v = np.linspace(1.0, 0.0, height, dtype=np.float32)
uu, vv = np.meshgrid(u, v)
vertex_uvs = np.stack((uu, vv), axis=-1).reshape(-1, 2)
pixel_intrinsics = intrinsics.copy()
pixel_intrinsics[0] *= width
pixel_intrinsics[1] *= height
pixel_intrinsics[2] = (0.0, 0.0, 1.0)
point_depth = depth.reshape(-1).astype(np.float32)
face_mask = mask.reshape(-1)[faces].all(axis=1)
face_depth_discontinuity = depth_discontinuity.reshape(-1)[faces].max(axis=1)
masked_depth = point_depth[mask.reshape(-1).astype(bool)]
reference_depth = np.median(masked_depth if len(masked_depth) else point_depth)
np.savez(
output_path,
vertices=vertices.astype(np.float32),
faces=faces.astype(np.int32),
uv=vertex_uvs.astype(np.float32),
intrinsics=pixel_intrinsics.astype(np.float32),
image_size=np.array((width, height), dtype=np.int32),
reference_depth=np.float32(reference_depth),
point_depth=point_depth,
normal_map=normal.astype(np.float32),
face_mask=face_mask,
face_depth_discontinuity=face_depth_discontinuity,
)
def generate_moge2(context, parameters):
"""Run MoGe-2 inference and return the artifact used by the requesting Job."""
import numpy as np
from PIL import Image
from .onnx_moge2 import (
infer as infer_moge2,
provider_label,
)
context.progress(0.02, "Starting MoGe-2")
output_directory = context.directory
output_directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
images, temporary_directory = prepare_input(
parameters["input"],
parameters.get("video_frames", 16),
output_directory,
start_frame=parameters.get("start_frame", 1),
cancel_check=context.check_cancelled,
)
try:
model = model_manager.get_moge2(model_directory, device)
context.progress(
0.08,
f"Loading MoGe-2 on {provider_label(model)}",
)
depths = []
normals = []
masks = []
points = []
intrinsics = []
source_images = []
for index, image_path in enumerate(images):
context.check_cancelled()
context.progress(
0.15 + 0.68 * index / len(images),
f"Estimating MoGe-2 geometry ({index + 1}/{len(images)})",
)
with Image.open(image_path) as opened:
image = np.asarray(opened.convert("RGB"), dtype=np.float32) / 255.0
source_images.append(image)
prediction = infer_moge2(
model,
image,
int(parameters["resolution_level"]),
)
depth = prediction["depth"]
normal = prediction["normal"]
mask = prediction["mask"].astype(bool)
point = prediction["points"]
intrinsic = prediction["intrinsics"]
depths.append(depth)
normals.append(normal)
masks.append(mask)
intrinsics.append(intrinsic)
if bool(
parameters.get("include_points", False)
or parameters.get("surface_mesh", False)
):
points.append(point)
debug_format = parameters.get("export_format", "")
if not debug_format or debug_format == "maps":
_save_moge_preview(
output_directory,
None if len(images) == 1 else index,
depth,
normal,
mask,
)
shapes = {array.shape for array in depths}
if len(shapes) != 1:
raise RuntimeError("MoGe-2 frame sizes must match")
result = {
"depth": np.stack(depths).astype(np.float32),
"normal": np.stack(normals).astype(np.float32),
"mask": np.stack(masks),
"intrinsics": np.stack(intrinsics).astype(np.float32),
}
if points:
result["points"] = np.stack(points).astype(np.float32)
debug_format = parameters.get("export_format", "")
if debug_format in {"glb", "ply"} and len(images) == 1:
context.progress(
0.84,
f"Exporting MoGe-2 {debug_format.upper()}",
)
export_moge_debug_mesh(
output_directory,
points[0],
depths[0],
masks[0],
normals[0],
source_images[0],
debug_format,
)
foreground_mask = None
remove_background_model = parameters.get("remove_background_model", "")
if remove_background_model:
remove_background_model_dir = model_manager.directory(
remove_background_model
)
context.progress( 0.87, "Removing image background")
if len(images) == 1:
foreground, _, _ = infer_foreground(
images[0],
remove_background_model_dir,
device,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
foreground_path = output_directory / "foreground.png"
foreground.save(foreground_path)
if points:
foreground_mask = foreground_mesh_mask(
foreground_path,
result["depth"][0].shape,
)
else:
infer_foreground_frames(
images,
remove_background_model_dir,
device,
output_directory,
context.progress,
0.87,
0.08,
model_cache=model_manager,
cancel_check=context.check_cancelled,
)
if bool(parameters.get("surface_mesh", False)):
if len(images) != 1 or not points:
raise RuntimeError("MoGe-2 Surface requires one point map")
context.progress( 0.94, "Building official MoGe-2 surface")
export_moge_surface_archive(
output_directory / "mesh.npz",
points[0],
result["depth"][0],
result["mask"][0],
result["normal"][0],
result["intrinsics"][0],
foreground_mask=(
foreground_mask
if remove_background_model
else None
),
)
np.savez(output_directory / "results.npz", **result)
copy_frames(images, output_directory)
context.progress(
1.0,
f"Exported MoGe-2 data for {len(images)} frame(s)",
)
if context.job_type == "generate-moge2-surface":
return {"surface": "mesh.npz"}
if context.job_type == "generate-moge2-depth":
return {"depth": "results.npz"}
return {"directory": "."}
finally:
if temporary_directory is not None:
shutil.rmtree(temporary_directory, ignore_errors=True)

1
src/snapform/ai_service/ben2_onnx.py → src/snapform/server/onnx_ben2.py
View File

@@ -48,3 +48,4 @@ def infer(session, image):
foreground = image.convert("RGBA")
foreground.putalpha(mask)
return foreground

1
src/snapform/ai_service/da3_onnx.py → src/snapform/server/onnx_da3.py
View File

@@ -240,3 +240,4 @@ def infer(
processed_images=processed_images,
aux={},
)

1
src/snapform/ai_service/moge_onnx.py → src/snapform/server/onnx_moge2.py
View File

@@ -163,3 +163,4 @@ def infer(session, image, resolution_level):
raw_output = dict(zip(names, values))
raw_output["metric_scale"] = raw_output.pop(scale_name)
return postprocess(raw_output)

0
src/snapform/ai_service/onnx_runtime.py → src/snapform/server/onnx_runtime.py
View File

55
src/snapform/status_bar.py
View File

@@ -1,55 +0,0 @@
from .operators import active_job
def draw_status_bar(self, context):
scene = getattr(context, "scene", None)
job = active_job()
if job is None or scene is None:
return
settings = getattr(scene, "snapform_settings", None)
if settings is None:
return
progress = min(max(settings.progress, 0.0), 1.0)
row = self.layout.row(align=True)
progress_row = row.row(align=True)
progress_row.ui_units_x = 18.0
progress_row.progress(
factor=progress,
type="BAR",
text=settings.status,
)
row.operator(
"snapform.cancel_job",
text="",
icon="X",
)
def tag_ui_regions(context):
window_manager = getattr(context, "window_manager", None)
if window_manager is None:
return
for window in window_manager.windows:
screen = getattr(window, "screen", None)
if screen is None:
continue
for area in screen.areas:
if area.type in {"STATUSBAR", "VIEW_3D"}:
area.tag_redraw()
for region in getattr(area, "regions", ()):
region.tag_redraw()
def redraw_ui(context, force=False):
tag_ui_regions(context)
if not force:
return
workspace = getattr(context, "workspace", None)
if workspace is not None:
try:
workspace.status_text_set_internal(None)
except (AttributeError, RuntimeError, TypeError):
pass