"""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", }