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

feat(node_assets): add image plane geometry nodes and validation scripts

Browse Source 
- Implemented `image_plane.py` to build the Image Plane Geometry Nodes group with customizable parameters for geometry, subdivide, and thickness.
- Created `image_plane_validation.py` to validate the functionality of the image plane, ensuring correct geometry generation and material assignment.
- Developed `material_layers.py` to construct shader node groups for image materials, including depth and normal handling.
- Authored `plane_layouts.py` to define node arrangements for image geometry assets.
- Introduced `shadeless.py` to create a shadeless shader node group that displays color without scene lighting while preserving shadows and reflections.
This commit is contained in:
马南
2026-09-08 23:24:27 +08:00
parent 196b964efc
commit 4c902ed7db
155 changed files with 9700 additions and 10393 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/anyimage/server/app.py
View File

@@ -2,18 +2,6 @@
from blendjob import JobServer
from .jobs.debug_predictions import (
run_ben2 as _debug_ben2,
)
from .jobs.debug_predictions import (
run_da3 as _debug_da3,
)
from .jobs.debug_predictions import (
run_moge2 as _debug_moge2,
)
from .jobs.debug_predictions import (
run_upscale as _debug_upscale,
)
from .jobs.cutout import run as _generate_cutout_artifacts
from .jobs.depth_plane import run as _generate_depth_plane_geometry
from .jobs.remove_background import run as _remove_background
@@ -58,23 +46,3 @@ def generate_depth_plane_geometry(context, parameters):
@server.job("generate-cutout-artifacts")
def generate_cutout_artifacts(context, parameters):
return _generate_cutout_artifacts(context, parameters)
@server.job("debug-da3")
def debug_da3(context, parameters):
return _debug_da3(context, parameters)
@server.job("debug-moge2")
def debug_moge2(context, parameters):
return _debug_moge2(context, parameters)
@server.job("debug-ben2")
def debug_ben2(context, parameters):
return _debug_ben2(context, parameters)
@server.job("debug-upscale")
def debug_upscale(context, parameters):
return _debug_upscale(context, parameters)

18
src/anyimage/server/geometry/depth_texture.py
View File

@@ -71,12 +71,12 @@ def _write_float_exr(rgba, output_path):
return output_path
def write_vector_depth_texture(frame, output_path):
def write_depth_texture(frame, output_path):
"""Write a frame's camera-space XYZ and validity as float RGBA."""
import numpy as np
if frame.points is None:
raise ValueError("Vector depth requires a point field")
raise ValueError("Depth requires a point field")
points = np.asarray(frame.points, dtype=np.float32)
validity = np.asarray(frame.validity, dtype=np.float32)
if not np.isfinite(points).all():
@@ -87,20 +87,6 @@ def write_vector_depth_texture(frame, output_path):
return _write_float_exr(rgba, output_path)
def write_z_depth_texture(frame, output_path):
"""Write a frame's camera Z and validity as float RGBA."""
import numpy as np
depth = np.asarray(frame.depth, dtype=np.float32)
validity = np.asarray(frame.validity, dtype=np.float32)
if not np.isfinite(depth).all():
raise ValueError("Depth field contains non-finite values")
rgba = np.empty((*depth.shape, 4), dtype=np.float32)
rgba[..., :3] = depth[..., None]
rgba[..., 3] = validity
return _write_float_exr(rgba, output_path)
def write_depth_metadata(frame, reference_valid, output_path):
"""Write camera values needed to scale a model-space depth texture."""
import numpy as np

2
src/anyimage/server/geometry/normal_texture.py
View File

@@ -65,5 +65,3 @@ def write_camera_normal_texture(
output_path = Path(output_path)
full.save(output_path)
return output_path

27
src/anyimage/server/geometry/prediction_artifacts.py
View File

@@ -4,8 +4,7 @@ from ..media import DEFAULT_MAX_AI_INPUT_SIZE, limit_image
from ..models import moge2
from .depth_texture import (
write_depth_metadata,
write_vector_depth_texture,
write_z_depth_texture,
write_depth_texture,
)
from .normal_texture import write_camera_normal_texture
from .edge_depth import extend_edge_depth
@@ -57,16 +56,13 @@ def generate_moge2_artifacts(
parameters,
input_path,
*,
depth_mode="NONE",
generate_depth=False,
normal_mode="NONE",
correct_edge_depth=False,
alpha_threshold=0.9,
):
"""Run MoGe-2 once and write only the requested file artifacts."""
depth_mode = str(depth_mode).upper()
normal_mode = str(normal_mode).upper()
if depth_mode not in {"NONE", "VECTOR", "Z"}:
raise ValueError(f"Unsupported depth mode: {depth_mode}")
if normal_mode not in {"NONE", "TANGENT", "OBJECT"}:
raise ValueError(f"Unsupported normal mode: {normal_mode}")
model_path = _limited_model_input(
@@ -79,32 +75,23 @@ def generate_moge2_artifacts(
context.resource("model_manager"),
parameters,
[model_path],
include_points=depth_mode == "VECTOR",
include_points=generate_depth,
cancel_check=context.check_cancelled,
)
if len(frames) != 1:
raise ValueError("Artifact generation requires one geometry frame")
frame = frames[0]
result = {}
if depth_mode != "NONE":
if generate_depth:
context.progress(0.45, "Generating depth texture")
alpha = _image_alpha(input_path, frame.depth.shape)
depth_frame = (
extend_edge_depth(frame, alpha, alpha_threshold)
if correct_edge_depth else frame
)
if depth_mode == "VECTOR":
depth_path = write_vector_depth_texture(
depth_frame,
context.directory / "vector-depth.exr",
)
result["vector_depth"] = depth_path
else:
depth_path = write_z_depth_texture(
depth_frame,
context.directory / "z-depth.exr",
)
result["z_depth"] = depth_path
result["depth"] = write_depth_texture(
depth_frame, context.directory / "depth.exr",
)
result["depth_metadata"] = write_depth_metadata(
depth_frame,
(alpha > 0.0) & (alpha >= alpha_threshold)

8
src/anyimage/server/jobs/cutout.py
View File

@@ -10,16 +10,16 @@ def run(context, parameters):
context.directory.mkdir(parents=True, exist_ok=True)
input_path = Path(parameters["input"])
result = {}
depth_mode = str(parameters.get("depth_mode", "NONE"))
generate_depth = bool(parameters.get("generate_depth", False))
normal_mode = str(parameters.get("normal_mode", "NONE"))
if depth_mode != "NONE" or normal_mode != "NONE":
if generate_depth or normal_mode != "NONE":
artifacts = generate_moge2_artifacts(
context,
parameters,
input_path,
depth_mode=depth_mode,
generate_depth=generate_depth,
normal_mode=normal_mode,
correct_edge_depth=depth_mode.upper() != "NONE",
correct_edge_depth=generate_depth,
alpha_threshold=float(parameters.get("alpha_threshold", 0.9)),
)
result.update({key: path.name for key, path in artifacts.items()})

96
src/anyimage/server/jobs/debug_predictions.py
View File

@@ -1,96 +0,0 @@
from ..media import copy_frames, input_frames
from ..models import da3
from ..models.ben2 import infer_foreground
from ..models.moge2 import infer_one as infer_moge2
from ..models.upscale import infer_one as infer_upscale
from ..outputs.debug_da3 import write_prediction as write_da3_prediction
from ..outputs.debug_moge2 import write_prediction as write_moge2_prediction
def run_da3(context, parameters):
context.progress(0.02, "Preparing DA3 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
prediction = da3.infer(
context.resource("model_manager"),
parameters,
images,
cancel_check=context.check_cancelled,
)
write_da3_prediction(prediction, context.directory / "prediction.npz")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw DA3 API prediction")
return {"directory": "."}
def run_moge2(context, parameters):
context.progress(0.02, "Preparing MoGe-2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_moge2(
context.resource("model_manager"),
parameters,
image,
release_memory=index == len(images) - 1,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
write_moge2_prediction(
prediction,
context.directory / f"prediction{suffix}.npz",
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw MoGe-2 API prediction")
return {"directory": "."}
def run_ben2(context, parameters):
context.progress(0.02, "Preparing BEN2 API input")
context.directory.mkdir(parents=True, exist_ok=True)
model_manager = context.resource("model_manager")
model_directory = model_manager.directory(parameters["model"])
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
foreground, _load_ms, _inference_ms = infer_foreground(
image,
model_directory,
parameters["device"],
model_cache=model_manager,
cancel_check=context.check_cancelled,
release_memory=index == len(images) - 1,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
foreground.save(context.directory / f"foreground{suffix}.png")
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw BEN2 API foreground")
return {"directory": "."}
def run_upscale(context, parameters):
context.progress(0.02, "Preparing Real-ESRGAN API input")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
for index, image in enumerate(images):
context.check_cancelled()
prediction = infer_upscale(
context.resource("model_manager"),
parameters,
image,
cancel_check=context.check_cancelled,
release_memory=index == len(images) - 1,
)
suffix = "" if len(images) == 1 else f".{index + 1:04d}"
prediction.save(
context.directory / f"prediction{suffix}.png",
format="PNG",
compress_level=1,
)
copy_frames(images, context.directory)
context.check_cancelled()
context.progress(1.0, "Saved raw Real-ESRGAN API prediction")
return {"directory": "."}

10
src/anyimage/server/jobs/depth_plane.py
View File

@@ -1,23 +1,23 @@
"""Generate model-space Z Depth and preserve color for one complete image."""
"""Generate model-space depth and preserve color for one complete image."""
from ..geometry.prediction_artifacts import generate_moge2_artifacts
from ..media import copy_color_images, input_frames
def run(context, parameters):
context.progress(0.05, "Preparing Depth Plane")
context.progress(0.05, "Preparing depth conversion")
context.directory.mkdir(parents=True, exist_ok=True)
with input_frames(context, parameters) as images:
if len(images) != 1:
raise ValueError("Depth Plane currently supports one image")
raise ValueError("Depth conversion currently supports one image")
copy_color_images(images, context.directory)
artifacts = generate_moge2_artifacts(
context,
parameters,
images[0],
depth_mode="Z",
generate_depth=True,
normal_mode="TANGENT",
)
context.check_cancelled()
context.progress(1.0, "Depth Plane data is ready")
context.progress(1.0, "Depth conversion data is ready")
return {key: path.name for key, path in artifacts.items()}

1
src/anyimage/server/jobs/upscale.py
View File

@@ -29,7 +29,6 @@ def run(context, parameters):
parameters,
image,
cancel_check=context.check_cancelled,
preserve_alpha=True,
release_memory=index == total - 1,
)
prediction.save(

2
src/anyimage/server/media/__init__.py
View File

@@ -6,7 +6,6 @@ from .input import (
check_cancelled,
collect_images,
copy_color_images,
copy_frames,
extract_video_frames,
first_frames,
input_frames,
@@ -31,7 +30,6 @@ __all__ = (
"check_cancelled",
"collect_images",
"copy_color_images",
"copy_frames",
"extract_video_frames",
"first_frames",
"input_frames",

13
src/anyimage/server/media/input.py
View File

@@ -157,16 +157,3 @@ def copy_color_images(images, output_directory):
if Path(image_path).resolve() == destination.resolve():
continue
shutil.copy2(image_path, destination)
def copy_frames(images, output_directory):
"""Copy multiple input images as numbered raw Debug frame files."""
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)

80
src/anyimage/server/model_catalog.py
View File

@@ -9,70 +9,13 @@ class ModelSpec:
repository: str
license_name: str
description: str
supports_gs: bool
huggingface_repository: str = ""
download_patterns: tuple[str, ...] = ()
ready_patterns: tuple[str, ...] = ("config.json", "*.safetensors")
enabled: bool = True
r2_directory: str = ""
r2_files: tuple[tuple[str, int, str], ...] = ()
MODELS = {
"SMALL": ModelSpec(
key="SMALL",
label="DA3 Small",
directory_name="DA3-SMALL",
repository="depth-anything/DA3-SMALL",
license_name="Apache 2.0",
huggingface_repository="depth-anything/DA3-SMALL",
description="Fastest any-view model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-SMALL",
r2_files=(
("model.onnx", 105484029, "7a9d81475b3e2334ace0d6d2c2a816b3c93081437e1866cc434ce1789125f631"),
),
),
"BASE": ModelSpec(
key="BASE",
label="DA3 Base",
directory_name="DA3-BASE",
repository="depth-anything/DA3-BASE",
license_name="Apache 2.0",
huggingface_repository="depth-anything/DA3-BASE",
description="Balanced any-view model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="DA3-BASE",
r2_files=(
("model.onnx", 412856697, "31f38fc4990539d7902ba267924b9731db1acb00d840c248a60b1865810ff9e6"),
),
),
"GIANT_1_1": ModelSpec(
key="GIANT_1_1",
label="DA3 Giant 1.1",
directory_name="da3_giant_1_1",
repository="depth-anything/DA3-GIANT-1.1",
license_name="CC BY-NC 4.0",
description="Non-commercial Giant model with Gaussian head",
supports_gs=True,
enabled=False,
),
"NESTED_GIANT_LARGE_1_1": ModelSpec(
key="NESTED_GIANT_LARGE_1_1",
label="DA3 Nested Giant Large 1.1",
directory_name="da3_nested_giant_large_1_1",
repository="depth-anything/DA3NESTED-GIANT-LARGE-1.1",
license_name="CC BY-NC 4.0",
description="Non-commercial nested model with metric depth and Gaussian head",
supports_gs=True,
enabled=False,
),
}
BEN2_MODEL = ModelSpec(
key="BEN2_BASE",
label="BEN2 Base",
@@ -81,7 +24,6 @@ BEN2_MODEL = ModelSpec(
license_name="Apache 2.0",
huggingface_repository="onnx-community/BEN2-ONNX",
description="Fast high-quality FP16 ONNX background removal model",
supports_gs=False,
download_patterns=("onnx/model_fp16.onnx",),
ready_patterns=("onnx/model_fp16.onnx",),
r2_directory="BEN2-ONNX",
@@ -100,7 +42,6 @@ REALESRGAN_GENERAL_WDN_X4V3 = ModelSpec(
repository="xinntao/Real-ESRGAN",
license_name="BSD 3-Clause",
description="Fast and conservative 2x output from a 4x ONNX model",
supports_gs=False,
download_patterns=("realesr-general-wdn-x4v3.onnx",),
ready_patterns=("realesr-general-wdn-x4v3.onnx",),
r2_directory="realesr-general-wdn-x4v3",
@@ -118,7 +59,6 @@ SAFMN_MODEL = ModelSpec(
repository="sunny2109/SAFMN",
license_name="Apache 2.0",
description="Balanced 2x output from a real-world 4x ONNX model",
supports_gs=False,
download_patterns=("safmn-real-x4.onnx",),
ready_patterns=("safmn-real-x4.onnx",),
r2_directory="safmn-real-x4",
@@ -137,7 +77,6 @@ REALESRGAN_X4PLUS = ModelSpec(
license_name="BSD 3-Clause",
huggingface_repository="jonathanst29/tinier-upscale-models",
description="Detailed 2x output from a general-purpose 4x ONNX model",
supports_gs=False,
download_patterns=("realesrgan-x4plus.onnx",),
ready_patterns=("realesrgan-x4plus.onnx",),
r2_directory="realesrgan-x4plus",
@@ -156,7 +95,6 @@ PISA_SR = ModelSpec(
repository="csslc/PiSA-SR",
license_name="Apache 2.0 / CreativeML Open RAIL++-M",
description="Generative 2x upscale with fixed detail strength 0.7",
supports_gs=False,
download_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
ready_patterns=("encoder.onnx", "decoder.onnx", "unet-pixel.onnx", "unet-full.onnx", "noise.npy", "LICENSE-APACHE.txt", "LICENSE-MODEL.txt", "NOTICE.txt"),
r2_directory="pisa-sr",
@@ -196,7 +134,6 @@ MOGE2_MODELS = {
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vits-normal-onnx",
description="Fastest MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vits-normal-onnx",
@@ -214,7 +151,6 @@ MOGE2_MODELS = {
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vitb-normal-onnx",
description="Balanced MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vitb-normal-onnx",
@@ -232,7 +168,6 @@ MOGE2_MODELS = {
license_name="MIT",
huggingface_repository="Ruicheng/moge-2-vitl-normal-onnx",
description="Highest-quality MoGe-2 model using ONNX Runtime",
supports_gs=False,
download_patterns=("model.onnx",),
ready_patterns=("model.onnx",),
r2_directory="moge-2-vitl-normal-onnx",
@@ -245,21 +180,10 @@ MOGE2_MODELS = {
}
MOGE2_MODEL = MOGE2_MODELS["MOGE2_VITS_NORMAL"]
DOWNLOADABLE_MODELS = {
**MODELS,
BEN2_MODEL.key: BEN2_MODEL,
**UPSCALE_MODELS,
**MOGE2_MODELS,
}
MODEL_NUMBERS = {
model_key: number for number, model_key in enumerate(MODELS)
}
def get_model(model_key):
try:
return MODELS[model_key]
except KeyError as error:
raise ValueError(f"Unknown DA3 model: {model_key}") from error
def get_downloadable_model(model_key):
@@ -281,16 +205,12 @@ def model_record(model, ready=False):
"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 "upscale"
if model.key in UPSCALE_MODELS
else "da3"
),
"ready": bool(ready),
}

25
src/anyimage/server/model_manager.py
View File

@@ -14,8 +14,6 @@ class ModelSessionCache:
def __init__(self):
self.ben2_model = None
self.ben2_key = None
self.da3_model = None
self.da3_key = None
self.moge_model = None
self.moge_key = None
self.upscale_model = None
@@ -39,21 +37,6 @@ class ModelSessionCache:
elapsed_ms = (time.perf_counter() - started) * 1000.0
return model, elapsed_ms
def get_da3(self, model_dir, device):
requested_device = getattr(device, "type", device)
key = (str(Path(model_dir).resolve()), requested_device, "onnx")
if self.da3_model is not None and self.da3_key == key:
return self.da3_model
self.release_geometry()
return self._load_da3(model_dir, requested_device, key)
def _load_da3(self, model_dir, requested_device, key):
from .models.onnx_da3 import create_session
model = create_session(model_dir, requested_device)
self.da3_model = model
self.da3_key = key
return model
def get_moge2(self, model_dir, requested_device):
from .models.onnx_moge2 import create_session
@@ -88,8 +71,6 @@ class ModelSessionCache:
self._release_runtime_cache()
def release_geometry(self):
self.da3_model = None
self.da3_key = None
self.moge_model = None
self.moge_key = None
self._release_runtime_cache()
@@ -107,17 +88,15 @@ class ModelSessionCache:
if close is not None:
close()
self.ben2_model = None
self.da3_model = None
self.moge_model = None
self.upscale_model = None
self.ben2_key = None
self.da3_key = None
self.moge_key = None
self.upscale_key = None
self._release_runtime_cache()
def snapshot(self):
geometry_key = self.da3_key or self.moge_key
geometry_key = self.moge_key
return {
"ben2": Path(self.ben2_key[0]).name if self.ben2_key else "",
"geometry": Path(geometry_key[0]).name if geometry_key else "",
@@ -169,8 +148,6 @@ class ModelManager:
def get_ben2(self, model_dir, device):
return self.sessions.get_ben2(model_dir, device)
def get_da3(self, model_dir, device):
return self.sessions.get_da3(model_dir, device)
def get_moge2(self, model_dir, device):
return self.sessions.get_moge2(model_dir, device)

3
src/anyimage/server/models/__init__.py
View File

@@ -1,8 +1,7 @@
from . import ben2, da3, moge2, upscale
from . import ben2, moge2, upscale
__all__ = (
"ben2",
"da3",
"moge2",
"upscale",
)

109
src/anyimage/server/models/da3.py
View File

@@ -1,109 +0,0 @@
from ..geometry import unproject_camera_depth
from .geometry import GeometryFrame
def merge_frame_predictions(predictions):
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",
):
setattr(
merged,
name,
np.concatenate(
[getattr(prediction, name) for prediction in predictions],
axis=0,
),
)
return merged
def infer(
model_manager,
parameters,
images,
*,
progress_callback=None,
cancel_check=None,
):
from .onnx_da3 import infer as infer_onnx
from .onnx_runtime import provider_label
model_directory = model_manager.directory(parameters["model"])
device = parameters["device"]
model = model_manager.get_da3(model_directory, device)
options = {
"process_res": parameters["process_res"],
"process_res_method": parameters["process_res_method"],
"ref_view_strategy": parameters["ref_view_strategy"],
}
def run(paths, *, release_memory):
try:
return infer_onnx(
model,
paths,
**options,
release_memory=release_memory,
)
except UnicodeDecodeError as error:
if provider_label(model) == "DirectML":
raise RuntimeError(
"DA3 inference failed on DirectML. Switch the device to "
"CPU in AnyImage Preferences and retry. "
f"Original error: {error}"
) from error
raise
if parameters.get("inference_mode", "joint") != "per_frame":
return run(
[str(path) for path in images],
release_memory=True,
)
predictions = []
for index, image in enumerate(images):
if cancel_check is not None:
cancel_check()
if progress_callback is not None:
progress_callback(index, len(images))
predictions.append(
run(
[str(image)],
release_memory=index == len(images) - 1,
)
)
return merge_frame_predictions(predictions)
def geometry_frames(prediction):
import numpy as np
depth = np.asarray(prediction.depth, dtype=np.float32)
validity = np.asarray(prediction.conf, dtype=np.float32)
intrinsics = np.asarray(prediction.intrinsics, dtype=np.float32)
if validity.shape != depth.shape or intrinsics.shape != (len(depth), 3, 3):
raise ValueError("DA3 prediction fields do not share one frame layout")
return [
GeometryFrame(
depth=frame_depth,
validity=frame_validity,
intrinsics=frame_intrinsics,
points=unproject_camera_depth(frame_depth, frame_intrinsics).astype(
np.float32
),
)
for frame_depth, frame_validity, frame_intrinsics in zip(
depth,
validity,
intrinsics,
)
]

1
src/anyimage/server/models/onnx_ben2.py
View File

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

246
src/anyimage/server/models/onnx_da3.py
View File

@@ -1,246 +0,0 @@
"""Depth Anything 3 inference with ONNX Runtime and NumPy."""
from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
import numpy as np
try:
from .onnx_runtime import create_session as create_onnx_session
from .onnx_runtime import run_session
except ImportError:
from onnx_runtime import create_session as create_onnx_session
from onnx_runtime import run_session
ONNX_PATH = "model.onnx"
PATCH_SIZE = 14
# DA3 cross-view attention is quadratic in (views x patches); the peak
# attention tensor alone can exhaust memory long before ONNX Runtime reports
# a generic allocation failure.
ATTENTION_BUDGET_BYTES = 4 * 1024 * 1024 * 1024
STRATEGY_CODES = {
"first": 0,
"middle": 1,
"saddle_balanced": 2,
"saddle_sim_range": 3,
}
IMAGE_MEAN = np.asarray((0.485, 0.456, 0.406), dtype=np.float32)
IMAGE_STD = np.asarray((0.229, 0.224, 0.225), dtype=np.float32)
@dataclass
class Prediction:
depth: np.ndarray
conf: np.ndarray
extrinsics: np.ndarray
intrinsics: np.ndarray
processed_images: np.ndarray
is_metric: int = 0
sky: np.ndarray | None = None
gaussians: object | None = None
aux: dict | None = None
scale_factor: float | None = None
def create_session(model_dir, requested_device="auto"):
path = Path(model_dir) / ONNX_PATH
return create_onnx_session(path, requested_device)
def nearest_multiple(value, multiple=PATCH_SIZE):
lower = value // multiple * multiple
upper = lower + multiple
return upper if abs(upper - value) <= abs(value - lower) else lower
def target_size(size, process_res, method):
width, height = size
boundary = max(width, height)
if method.startswith("lower_bound"):
boundary = min(width, height)
scale = float(process_res) / boundary
width = max(1, int(round(width * scale)))
height = max(1, int(round(height * scale)))
if method.endswith("resize"):
return (
max(PATCH_SIZE, nearest_multiple(width)),
max(PATCH_SIZE, nearest_multiple(height)),
)
if method.endswith("crop"):
return (
max(PATCH_SIZE, width // PATCH_SIZE * PATCH_SIZE),
max(PATCH_SIZE, height // PATCH_SIZE * PATCH_SIZE),
)
raise ValueError(f"Unsupported process_res_method: {method}")
def resize_and_crop(image, process_res, method):
from PIL import Image
width, height = image.size
boundary = max(width, height)
if method.startswith("lower_bound"):
boundary = min(width, height)
scale = float(process_res) / boundary
resized_size = (
max(1, int(round(width * scale))),
max(1, int(round(height * scale))),
)
resampling = (
Image.Resampling.BICUBIC
if scale > 1.0
else Image.Resampling.BOX
)
image = image.resize(resized_size, resample=resampling)
final_size = target_size((width, height), process_res, method)
if method.endswith("resize"):
if image.size != final_size:
upscale = final_size[0] > image.width or final_size[1] > image.height
resampling = (
Image.Resampling.BICUBIC
if upscale
else Image.Resampling.BOX
)
image = image.resize(final_size, resample=resampling)
return image
left = (image.width - final_size[0]) // 2
top = (image.height - final_size[1]) // 2
return image.crop((left, top, left + final_size[0], top + final_size[1]))
def preprocess(image_paths, process_res=504, process_res_method="upper_bound_resize"):
from PIL import Image
processed = []
for path in image_paths:
with Image.open(path) as source:
image = resize_and_crop(
source.convert("RGB"),
process_res,
process_res_method,
)
processed.append(np.asarray(image, dtype=np.uint8))
min_height = min(image.shape[0] for image in processed)
min_width = min(image.shape[1] for image in processed)
cropped = []
for image in processed:
top = (image.shape[0] - min_height) // 2
left = (image.shape[1] - min_width) // 2
cropped.append(
image[top : top + min_height, left : left + min_width]
)
images = np.stack(cropped)
normalized = images.astype(np.float32) / 255.0
normalized = (normalized - IMAGE_MEAN) / IMAGE_STD
model_input = normalized.transpose(0, 3, 1, 2)[None]
return np.ascontiguousarray(model_input), images
def quaternion_to_matrix(quaternion):
i, j, k, real = np.moveaxis(quaternion, -1, 0)
two_scale = 2.0 / np.sum(quaternion * quaternion, axis=-1)
values = (
1 - two_scale * (j * j + k * k),
two_scale * (i * j - k * real),
two_scale * (i * k + j * real),
two_scale * (i * j + k * real),
1 - two_scale * (i * i + k * k),
two_scale * (j * k - i * real),
two_scale * (i * k - j * real),
two_scale * (j * k + i * real),
1 - two_scale * (i * i + j * j),
)
return np.stack(values, axis=-1).reshape(quaternion.shape[:-1] + (3, 3))
def camera_from_pose(pose_encoding, height, width):
translation = pose_encoding[..., :3]
rotation = quaternion_to_matrix(pose_encoding[..., 3:7])
rotation_inverse = np.swapaxes(rotation, -1, -2)
translation_inverse = -rotation_inverse @ translation[..., None]
extrinsics = np.concatenate(
(rotation_inverse, translation_inverse),
axis=-1,
).astype(np.float32)
field_of_view_height = pose_encoding[..., 7]
field_of_view_width = pose_encoding[..., 8]
focal_y = (height / 2.0) / np.maximum(
np.tan(field_of_view_height / 2.0),
1e-6,
)
focal_x = (width / 2.0) / np.maximum(
np.tan(field_of_view_width / 2.0),
1e-6,
)
intrinsics = np.zeros(pose_encoding.shape[:-1] + (3, 3), dtype=np.float32)
intrinsics[..., 0, 0] = focal_x
intrinsics[..., 1, 1] = focal_y
intrinsics[..., 0, 2] = width / 2.0
intrinsics[..., 1, 2] = height / 2.0
intrinsics[..., 2, 2] = 1.0
return extrinsics, intrinsics
def estimate_attention_bytes(views, height, width):
"""Estimate the peak cross-view attention tensor size in bytes."""
patches = (height // PATCH_SIZE) * (width // PATCH_SIZE)
tokens = views * patches
# 12 attention heads in DA3 Base; float32 activations.
return tokens * tokens * 12 * 4
def infer(
session,
image_paths,
process_res=504,
process_res_method="upper_bound_resize",
ref_view_strategy="saddle_balanced",
*,
release_memory=True,
):
inputs = session.get_inputs()
model_input, processed_images = preprocess(
image_paths,
process_res,
process_res_method,
)
_, views, _, height, width = model_input.shape
attention_bytes = estimate_attention_bytes(views, height, width)
if attention_bytes > ATTENTION_BUDGET_BYTES:
raise RuntimeError(
f"DA3 cannot process {views} frame(s) at {width}x{height}: "
"the cross-view attention would need about "
f"{attention_bytes // 2**30} GB. Reduce the frame count or "
"the Resolution Scale."
)
strategy_code = np.asarray(
[STRATEGY_CODES.get(ref_view_strategy, 2)],
dtype=np.int64,
)
input_names = [value.name for value in inputs]
outputs = run_session(
session,
None,
{
input_names[0]: model_input,
input_names[1]: strategy_code,
},
release_memory=release_memory,
)
depth, confidence, pose_encoding = outputs[:3]
height, width = depth.shape[-2:]
extrinsics, intrinsics = camera_from_pose(
pose_encoding[0],
height,
width,
)
return Prediction(
depth=np.asarray(depth[0], dtype=np.float32),
conf=np.asarray(confidence[0], dtype=np.float32),
extrinsics=extrinsics,
intrinsics=intrinsics,
processed_images=processed_images,
aux={},
)

2
src/anyimage/server/models/onnx_upscale.py
View File

@@ -153,5 +153,3 @@ def infer(session, image, cancel_check=None, *, release_memory=True):
raise UpscaleResourceError(
"the complete upscale result exceeded available system memory"
) from error

3
src/anyimage/server/models/upscale.py
View File

@@ -12,7 +12,6 @@ def infer_one(
parameters,
image_path,
cancel_check=None,
preserve_alpha=False,
release_memory=True,
):
from PIL import Image
@@ -33,7 +32,7 @@ def infer_one(
opened,
parameters.get("max_input_size", DEFAULT_MAX_AI_INPUT_SIZE),
)
source_alpha = _source_alpha(source) if preserve_alpha else None
source_alpha = _source_alpha(source)
prediction = infer(
model,
source,

1
src/anyimage/server/outputs/__init__.py
View File

@@ -1 +0,0 @@
"""Debug output serializers used by Job Server workflows."""

24
src/anyimage/server/outputs/debug_da3.py
View File

@@ -1,24 +0,0 @@
from pathlib import Path
def write_prediction(prediction, output_path):
import numpy as np
values = {}
for name in (
"depth",
"conf",
"extrinsics",
"intrinsics",
"processed_images",
"is_metric",
"sky",
"scale_factor",
):
value = getattr(prediction, name, None)
if value is not None:
values[name] = np.asarray(value)
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(output_path, **values)
return output_path

13
src/anyimage/server/outputs/debug_moge2.py
View File

@@ -1,13 +0,0 @@
from pathlib import Path
def write_prediction(prediction, output_path):
import numpy as np
output_path = Path(output_path)
output_path.parent.mkdir(parents=True, exist_ok=True)
np.savez(
output_path,
**{name: np.asarray(value) for name, value in prediction.items()},
)
return output_path