snomiaoandClaude Opus 4.6
dc2c298ce7
Two nodes: InfiniDepthModelLoader and InfiniDepthEstimate. Supports CUDA, MPS, and CPU via ComfyUI device management. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
271 lines
10 KiB
Python
271 lines
10 KiB
Python
"""ComfyUI nodes for InfiniDepth depth estimation."""
|
|
|
|
import os
|
|
import torch
|
|
import numpy as np
|
|
import cv2
|
|
|
|
import folder_paths
|
|
import comfy.model_management as mm
|
|
import comfy.utils
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Model loader node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class InfiniDepthModelLoader:
|
|
"""Load an InfiniDepth depth-estimation model from a checkpoint file."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_type": (["InfiniDepth", "InfiniDepth_DepthSensor"],),
|
|
"depth_ckpt": ("STRING", {
|
|
"default": "checkpoints/depth/infinidepth.ckpt",
|
|
"tooltip": "Path to depth checkpoint (relative to InfiniDepth repo root, or absolute).",
|
|
}),
|
|
"moge2_ckpt": ("STRING", {
|
|
"default": "checkpoints/moge-2-vitl-normal/model.pt",
|
|
"tooltip": "Path to MoGe-2 checkpoint for metric scale recovery.",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("INFINIDEPTH_MODEL",)
|
|
RETURN_NAMES = ("model",)
|
|
FUNCTION = "load_model"
|
|
CATEGORY = "depth/InfiniDepth"
|
|
DESCRIPTION = "Load an InfiniDepth depth estimation model."
|
|
|
|
def load_model(self, model_type: str, depth_ckpt: str, moge2_ckpt: str):
|
|
from InfiniDepth.model import MODEL_REGISTRY
|
|
import InfiniDepth.model.model as _model_mod
|
|
|
|
# Resolve relative paths against the InfiniDepth repo root
|
|
repo_root = os.environ.get(
|
|
"INFINIDEPTH_ROOT",
|
|
os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..", "..", "..")),
|
|
)
|
|
if not os.path.isabs(depth_ckpt):
|
|
depth_ckpt = os.path.join(repo_root, depth_ckpt)
|
|
if not os.path.isabs(moge2_ckpt):
|
|
moge2_ckpt = os.path.join(repo_root, moge2_ckpt)
|
|
|
|
# Patch: the original model.__init__ calls self.cuda() and checks for CUDA.
|
|
# We override to let ComfyUI manage device placement instead.
|
|
_orig_init = _model_mod._BaseInfiniDepthModel.__init__
|
|
|
|
def _patched_init(self_model, *args, **kwargs):
|
|
# Temporarily make torch.cuda.is_available return True so __init__ passes
|
|
_real_avail = torch.cuda.is_available
|
|
torch.cuda.is_available = lambda: True
|
|
# Replace self.cuda() with a no-op during init
|
|
_real_cuda = self_model.__class__.cuda
|
|
self_model.__class__.cuda = lambda s: s
|
|
try:
|
|
_orig_init(self_model, *args, **kwargs)
|
|
finally:
|
|
torch.cuda.is_available = _real_avail
|
|
self_model.__class__.cuda = _real_cuda
|
|
|
|
_model_mod._BaseInfiniDepthModel.__init__ = _patched_init
|
|
try:
|
|
if model_type not in MODEL_REGISTRY:
|
|
raise ValueError(f"Unknown model type: {model_type}, options: {list(MODEL_REGISTRY.keys())}")
|
|
model = MODEL_REGISTRY[model_type](model_path=depth_ckpt)
|
|
finally:
|
|
_model_mod._BaseInfiniDepthModel.__init__ = _orig_init
|
|
|
|
model.eval()
|
|
print(f"[InfiniDepth] Loaded {model_type} from {depth_ckpt}")
|
|
|
|
payload = {
|
|
"model": model,
|
|
"model_type": model_type,
|
|
"moge2_ckpt": moge2_ckpt,
|
|
}
|
|
return (payload,)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Depth estimation node
|
|
# ---------------------------------------------------------------------------
|
|
|
|
class InfiniDepthEstimate:
|
|
"""Run InfiniDepth depth estimation on an image (or batch of images)."""
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model": ("INFINIDEPTH_MODEL",),
|
|
"image": ("IMAGE",),
|
|
"output_resolution_mode": (["upsample", "original", "specific"],),
|
|
"upsample_ratio": ("INT", {
|
|
"default": 1,
|
|
"min": 1,
|
|
"max": 8,
|
|
"step": 1,
|
|
"tooltip": "Scale factor when output_resolution_mode is 'upsample'.",
|
|
}),
|
|
"output_height": ("INT", {
|
|
"default": 768,
|
|
"min": 64,
|
|
"max": 8192,
|
|
"step": 8,
|
|
"tooltip": "Height when output_resolution_mode is 'specific'.",
|
|
}),
|
|
"output_width": ("INT", {
|
|
"default": 1024,
|
|
"min": 64,
|
|
"max": 8192,
|
|
"step": 8,
|
|
"tooltip": "Width when output_resolution_mode is 'specific'.",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "IMAGE",)
|
|
RETURN_NAMES = ("depth_raw", "depth_colored",)
|
|
FUNCTION = "estimate_depth"
|
|
CATEGORY = "depth/InfiniDepth"
|
|
DESCRIPTION = "Estimate depth from an image using InfiniDepth."
|
|
|
|
@torch.no_grad()
|
|
def estimate_depth(
|
|
self,
|
|
model,
|
|
image: torch.Tensor,
|
|
output_resolution_mode: str,
|
|
upsample_ratio: int,
|
|
output_height: int,
|
|
output_width: int,
|
|
):
|
|
from InfiniDepth.utils.inference_utils import (
|
|
prepare_metric_depth_inputs,
|
|
resolve_camera_intrinsics_for_inference,
|
|
build_scaled_intrinsics_matrix,
|
|
resolve_output_size_from_mode,
|
|
apply_sky_mask_to_depth,
|
|
)
|
|
from InfiniDepth.utils.sampling_utils import SAMPLING_METHODS
|
|
from InfiniDepth.utils.io_utils import depth_to_disparity
|
|
from InfiniDepth.utils.vis_utils import colorize_depth_maps
|
|
|
|
depth_model = model["model"]
|
|
model_type = model["model_type"]
|
|
moge2_ckpt = model["moge2_ckpt"]
|
|
|
|
device = mm.get_torch_device()
|
|
depth_model.to(device)
|
|
|
|
# ComfyUI IMAGE: [B, H, W, C] float32 0-1
|
|
B, org_h, org_w, C = image.shape
|
|
|
|
input_size = (768, 1024)
|
|
|
|
pbar = comfy.utils.ProgressBar(B)
|
|
all_depth_raw = []
|
|
all_depth_colored = []
|
|
|
|
for i in range(B):
|
|
# Convert from ComfyUI BHWC to model's BCHW (RGB only, strip alpha if present)
|
|
img_frame = image[i].cpu().numpy()[:, :, :3] # keep only RGB channels
|
|
img_np = (img_frame * 255).astype(np.uint8)
|
|
img_resized = cv2.resize(img_np, (input_size[1], input_size[0]), interpolation=cv2.INTER_AREA)
|
|
img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).unsqueeze(0).float() / 255.0
|
|
img_tensor = img_tensor.to(device)
|
|
|
|
# Prepare metric depth (MoGe-2 fallback)
|
|
gt_depth, prompt_depth, gt_depth_mask, use_gt_depth, moge2_intrinsics = prepare_metric_depth_inputs(
|
|
input_depth_path=None,
|
|
input_size=input_size,
|
|
image=img_tensor,
|
|
device=device,
|
|
moge2_pretrained=moge2_ckpt,
|
|
)
|
|
|
|
fx_org, fy_org, cx_org, cy_org, _ = resolve_camera_intrinsics_for_inference(
|
|
fx_org=None, fy_org=None, cx_org=None, cy_org=None,
|
|
org_h=org_h, org_w=org_w,
|
|
image=img_tensor,
|
|
moge2_pretrained=moge2_ckpt,
|
|
moge2_intrinsics=moge2_intrinsics,
|
|
)
|
|
|
|
gt = depth_to_disparity(gt_depth)
|
|
prompt = depth_to_disparity(prompt_depth)
|
|
|
|
_, _, h, w = img_tensor.shape
|
|
|
|
h_sample, w_sample = resolve_output_size_from_mode(
|
|
output_resolution_mode=output_resolution_mode,
|
|
org_h=org_h, org_w=org_w,
|
|
h=h, w=w,
|
|
output_size=(output_height, output_width),
|
|
upsample_ratio=upsample_ratio,
|
|
)
|
|
|
|
query_coord = SAMPLING_METHODS["2d_uniform"]((h_sample, w_sample)).unsqueeze(0).to(device)
|
|
pred_depth, _ = depth_model.inference(
|
|
image=img_tensor,
|
|
query_coord=query_coord,
|
|
gt_depth=gt,
|
|
gt_depth_mask=gt_depth_mask,
|
|
prompt_depth=prompt,
|
|
prompt_mask=prompt > 0,
|
|
)
|
|
pred_depthmap = pred_depth.permute(0, 2, 1).view(1, 1, h_sample, w_sample)
|
|
|
|
# Normalize depth to 0-1 for output
|
|
depth_np = pred_depthmap.squeeze().cpu().numpy()
|
|
d_min, d_max = depth_np.min(), depth_np.max()
|
|
if d_max - d_min > 1e-6:
|
|
depth_norm = (depth_np - d_min) / (d_max - d_min)
|
|
else:
|
|
depth_norm = np.zeros_like(depth_np)
|
|
|
|
# Raw depth as grayscale IMAGE [H, W, 3]
|
|
depth_gray = np.stack([depth_norm] * 3, axis=-1).astype(np.float32)
|
|
all_depth_raw.append(torch.from_numpy(depth_gray))
|
|
|
|
# Colorized depth
|
|
colored = colorize_depth_maps(depth_np, d_min, d_max, cmap="Spectral")
|
|
if isinstance(colored, torch.Tensor):
|
|
colored = colored.cpu().numpy()
|
|
colored = colored.squeeze()
|
|
if colored.ndim == 3 and colored.shape[0] in (3, 4):
|
|
colored = colored.transpose(1, 2, 0)
|
|
colored = colored[:, :, :3].astype(np.float32)
|
|
if colored.max() > 1.0:
|
|
colored = colored / 255.0
|
|
all_depth_colored.append(torch.from_numpy(colored))
|
|
|
|
pbar.update(1)
|
|
|
|
# Offload model
|
|
depth_model.to(mm.unet_offload_device())
|
|
mm.soft_empty_cache()
|
|
|
|
depth_raw = torch.stack(all_depth_raw, dim=0)
|
|
depth_colored = torch.stack(all_depth_colored, dim=0)
|
|
|
|
return (depth_raw, depth_colored,)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Registration
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"InfiniDepthModelLoader": InfiniDepthModelLoader,
|
|
"InfiniDepthEstimate": InfiniDepthEstimate,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"InfiniDepthModelLoader": "Load InfiniDepth Model",
|
|
"InfiniDepthEstimate": "InfiniDepth Depth Estimation",
|
|
}
|