"""Camera-space surface normals from a (B, H, W, 3) point map (v1 fallback)."""

finite = torch.isfinite(points).all(dim=-1)

pts = torch.where(finite.unsqueeze(-1), points, torch.zeros_like(points))

dx = pts[..., :, 2:, :] - pts[..., :, :-2, :]

dy = pts[..., 2:, :, :] - pts[..., :-2, :, :]

dx = torch.nn.functional.pad(dx.permute(0, 3, 1, 2), (1, 1, 0, 0)).permute(0, 2, 3, 1)

dy = torch.nn.functional.pad(dy.permute(0, 3, 1, 2), (0, 0, 1, 1)).permute(0, 2, 3, 1)

# dy x dx (not dx x dy) so the result is outward-facing in OpenCV (Y-down flips the right-hand rule), matching v2's predicted normals.

n = torch.cross(dy, dx, dim=-1)

n = torch.nn.functional.normalize(n, dim=-1)

"""Per-batch normalize 1/depth to [0, 1] using 0.1/99.9 percentile clipping."""

out = torch.zeros_like(depth)

for i in range(depth.shape[0]):

d = depth[i]

valid = torch.isfinite(d) & (d > 0)

if not valid.any():

continue

disp = torch.where(valid, 1.0 / d.clamp_min(1e-6), torch.zeros_like(d))

disp_valid = disp[valid]

lo = torch.quantile(disp_valid, 0.001)

hi = torch.quantile(disp_valid, 0.999)

scale = (hi - lo).clamp_min(1e-6)

norm = ((disp - lo) / scale).clamp(0.0, 1.0)

out[i] = torch.where(valid, norm, torch.zeros_like(norm))

@classmethod

def define_schema(cls):

return io.Schema(

node_id="LoadMoGeModel",

display_name="Load MoGe Model",

category="model/loaders",

inputs=[

io.Combo.Input("model_name", options=folder_paths.get_filename_list("geometry_estimation")),

],

outputs=[MoGeModelType.Output()],

)

@classmethod

def execute(cls, model_name) -> io.NodeOutput:

path = folder_paths.get_full_path_or_raise("geometry_estimation", model_name)

sd = comfy.utils.load_torch_file(path, safe_load=True)

"""Equirectangular panorama inference: split into 12 perspective views, run

@classmethod

def define_schema(cls):

return io.Schema(

node_id="MoGePanoramaInference",

search_aliases=["moge", "panorama", "depth", "geometry", "depth estimation", "geometry estimation"],

display_name="Run MoGe Panorama Inference",

category="image/geometry estimation",

description="Run MoGe on an equirectangular panorama by splitting it into 12 perspective views, running inference on each, and merging the results into a single depth map.",

inputs=[

MoGeModelType.Input("moge_model"),

io.Image.Input("image", tooltip="Equirectangular panorama (any aspect)."),

io.Int.Input("resolution_level", default=9, min=0, max=9,

tooltip="Per-view detail (0 = fastest, 9 = most detailed)."),

io.Int.Input("split_resolution", default=512, min=256, max=1024,

tooltip="Resolution of each perspective split."),

io.Int.Input("merge_resolution", default=1920, min=256, max=8192,

tooltip="Long-side resolution of the merged equirect distance map."),

io.Int.Input("batch_size", default=4, min=1, max=12,

tooltip="Views per inference batch (12 splits total)."),

io.Int.Input("refine_steps", default=3, min=0, max=8, advanced=True,

tooltip="MoGe-3 only: sparse volumetric refinement passes over the predicted depth. More passes sharpen fine detail and edges at a roughly linear cost. 0 disables refinement. Ignored by MoGe-1 / MoGe-2."),

],

outputs=[MoGeGeometry.Output(display_name="moge_geometry")],

)

@classmethod

def execute(cls, moge_model, image, resolution_level, split_resolution, merge_resolution, batch_size, refine_steps) -> io.NodeOutput:

if image.shape[0] != 1:

raise ValueError(f"MoGePanoramaInference takes a single image (got batch of {image.shape[0]})")

image = image[..., :3]

H, W = int(image.shape[1]), int(image.shape[2])

scale = min(merge_resolution / max(H, W), 1.0)

merge_h, merge_w = max(int(H * scale), 32), max(int(W * scale), 32)

extrinsics, intrinsics = get_panorama_cameras()

comfy.model_management.load_model_gpu(moge_model.patcher)

device = moge_model.load_device

img_chw = image[0].movedim(-1, -3).to(device=device, dtype=moge_model.dtype)

splits = split_panorama_image(img_chw, extrinsics, intrinsics, split_resolution)

n_views = splits.shape[0]

# Weight each lsmr solve by 4^level so the final-resolution solve doesn't leave the bar idle.

merge_levels: list[tuple[int, int]] = []

w_, h_ = merge_w, merge_h

while True:

merge_levels.append((w_, h_))

if max(w_, h_) <= 256:

break

w_, h_ = w_ // 2, h_ // 2

merge_levels.reverse()

solve_weight = {wh: 4 ** i for i, wh in enumerate(merge_levels)}

n_merge_view_units = n_views * len(merge_levels)

n_merge_solve_units = sum(solve_weight.values())

pbar = comfy.utils.ProgressBar(n_views + n_merge_view_units + n_merge_solve_units)

done = 0

distance_maps: list = []

masks: list = []

with tqdm(total=n_views, desc="MoGe panorama inference") as tq:

for i in range(0, n_views, batch_size):

batch = splits[i:i + batch_size]

# apply_metric_scale=False: per-view scales would not align across overlap seams.

result = moge_model.infer(batch, resolution_level=resolution_level,

fov_x=90.0, force_projection=True,

apply_mask=False, apply_metric_scale=False,

refine_steps=refine_steps)

distance_maps.extend(list(result["points"].float().norm(dim=-1).cpu().numpy()))

masks.extend(list(result["mask"].cpu().numpy()))

n = batch.shape[0]

done += n

pbar.update_absolute(done)

tq.update(n)

with tqdm(total=n_merge_view_units + n_merge_solve_units, desc="MoGe panorama merge: views") as tq:

def _on_merge_view():

nonlocal done

done += 1

pbar.update_absolute(done)

tq.update(1)

def _on_solve_start(w, h):

tq.set_description(f"MoGe panorama merge: solving {w}x{h}")

def _on_solve_end(w, h):

nonlocal done

weight = solve_weight[(w, h)]

done += weight

pbar.update_absolute(done)

tq.update(weight)

tq.set_description("MoGe panorama merge: views")

pano_depth, pano_mask = merge_panorama_depth(

merge_w, merge_h, distance_maps, masks, list(extrinsics), intrinsics,

on_view=_on_merge_view, on_solve_start=_on_solve_start, on_solve_end=_on_solve_end)

pano_depth = torch.from_numpy(pano_depth)

pano_mask = torch.from_numpy(pano_mask)

if (merge_h, merge_w) != (H, W):

pano_depth = torch.nn.functional.interpolate(pano_depth[None, None], size=(H, W), mode="bilinear", align_corners=False).squeeze()

pano_mask = torch.nn.functional.interpolate(pano_mask[None, None].float(), size=(H, W), mode="nearest").squeeze() > 0

# Pixels uncovered by any view's predicted foreground are unconstrained in the lsmr solve and stay at log_depth=0 (depth=1)

if pano_mask.any() and not pano_mask.all():

far = torch.quantile(pano_depth[pano_mask], 0.95) * 5.0

pano_depth = torch.where(pano_mask, pano_depth, far)

directions = torch.from_numpy(spherical_uv_to_directions(_uv_grid(H, W)))

points = (directions * pano_depth[..., None]).unsqueeze(0)

depth = pano_depth.unsqueeze(0)

mask = pano_mask.unsqueeze(0)

# Points stay in MoGe spherical coords; MoGePointMapToMesh applies the spherical->glTF rotation after triangulation

moge_geometry = {"points": points, "depth": depth, "mask": mask, "image": image.cpu()}

@classmethod

def define_schema(cls):

return io.Schema(

node_id="MoGeInference",

search_aliases=["moge", "depth", "geometry", "depth estimation", "geometry estimation"],

display_name="Run MoGe Inference",

description="Run MoGe on a single image to estimate depth and geometry.",

category="image/geometry estimation",

inputs=[

MoGeModelType.Input("moge_model"),

io.Image.Input("image"),

io.Int.Input("resolution_level", default=9, min=0, max=9,

tooltip="0 = fastest, 9 = most detail."),

io.Float.Input("fov_x_degrees", default=0.0, min=0.0, max=170.0, step=0.1, advanced=True,

tooltip="Horizontal field of view of the source camera. Sets the focal length used to unproject the depth map into 3D. 0 = auto-recover from the predicted points."),

io.Int.Input("batch_size", default=4, min=1, max=64,

tooltip="Images per inference call. Lower if you OOM on a long video / image set."),

io.Boolean.Input("force_projection", default=True, advanced=True),

io.Boolean.Input("apply_mask", default=True, advanced=True,

tooltip="Set masked-out (sky / invalid) pixels to inf in points and depth so meshing culls them. Disable to keep the raw predicted geometry everywhere; the mask is still returned separately."),

io.Int.Input("refine_steps", default=3, min=0, max=8, advanced=True,

tooltip="MoGe-3 only: sparse volumetric refinement passes over the predicted depth. More passes sharpen fine detail and edges at a roughly linear cost. 0 disables refinement. Ignored by MoGe-1 / MoGe-2."),

],

outputs=[MoGeGeometry.Output(display_name="moge_geometry")],

)

@classmethod

def execute(cls, moge_model, image, resolution_level, fov_x_degrees, batch_size, force_projection, apply_mask, refine_steps) -> io.NodeOutput:

image = image[..., :3]

bchw = image.movedim(-1, -3).contiguous()

B = bchw.shape[0]

fov = None if fov_x_degrees <= 0 else float(fov_x_degrees)

pbar = comfy.utils.ProgressBar(B)

chunks: list[dict] = []

with tqdm(total=B, desc="MoGe inference") as tq:

for i in range(0, B, batch_size):

chunk = bchw[i:i + batch_size]

chunks.append(moge_model.infer(chunk, resolution_level=resolution_level, fov_x=fov,

force_projection=force_projection, apply_mask=apply_mask,

refine_steps=refine_steps))

pbar.update_absolute(min(i + batch_size, B))

tq.update(chunk.shape[0])

def stack(field):

vals = [c[field] for c in chunks if field in c]

return torch.cat(vals, dim=0) if vals else None

moge_geometry = {"image": image.cpu()}

for field in ("points", "depth", "intrinsics", "mask", "normal"):

v = stack(field)

if v is not None:

moge_geometry[field] = v

"""Render a visualization or mask from a MOGE_GEOMETRY packet."""

@classmethod

def define_schema(cls):

return io.Schema(

node_id="MoGeRender",

search_aliases=["moge", "render", "geometry", "depth", "normal"],

display_name="Render MoGe Geometry",

description="Render a depth map or normal map from geometry data",

category="image/geometry estimation",

inputs=[

MoGeGeometry.Input("moge_geometry"),

io.Combo.Input("output", options=["depth", "depth_colored", "normal_opengl", "normal_directx", "mask"], default="depth",

tooltip="DirectX vs OpenGL controls the normal-map green-channel convention. DirectX: green = -Y down (Unreal). OpenGL: green = +Y up (Blender, Substance, Unity, glTF)."),

],

outputs=[io.Image.Output()],

)

@classmethod

def execute(cls, moge_geometry, output) -> io.NodeOutput:

is_normal = output in ("normal_directx", "normal_opengl")

opengl = output.endswith("_opengl")

# Pick the input tensor for the chosen mode and validate availability.

if output in ("depth", "depth_colored"):

if "depth" not in moge_geometry:

raise ValueError("moge_geometry has no depth output.")

src = moge_geometry["depth"]

elif is_normal:

if "normal" in moge_geometry:

src = moge_geometry["normal"]

elif "points" in moge_geometry:

src = moge_geometry["points"]

else:

raise ValueError("moge_geometry has neither normals nor points to derive normals from.")

elif output == "mask":

if "mask" not in moge_geometry:

raise ValueError("moge_geometry has no mask output.")

src = moge_geometry["mask"]

else:

raise ValueError(f"Unknown output mode: {output}")

B = src.shape[0]

pbar = comfy.utils.ProgressBar(B)

out: list[torch.Tensor] = []

with tqdm(total=B, desc=f"MoGe render: {output}") as tq:

for i in range(B):

slc = src[i:i + 1].float()

if output in ("depth", "depth_colored"):

d = _normalize_disparity(slc)

out.append(_turbo(d) if output == "depth_colored"

else d.unsqueeze(-1).expand(*d.shape, 3).contiguous())

elif is_normal:

n = slc if "normal" in moge_geometry else _normals_from_points(slc)

# MoGe is OpenCV (Z+ into scene); normal-map convention is Z+ out of surface, so flip Z.

y_sign = -1.0 if opengl else 1.0

n = n * n.new_tensor([1.0, y_sign, -1.0])

out.append((n * 0.5 + 0.5).clamp(0.0, 1.0))

elif output == "mask":

out.append(slc.unsqueeze(-1).expand(*slc.shape, 3).contiguous())

pbar.update_absolute(i + 1)

tq.update(1)

result = torch.cat(out, dim=0).to(device=comfy.model_management.intermediate_device(), dtype=comfy.model_management.intermediate_dtype())

"""Triangulate one image of a MoGe point map into a Types.MESH (UVs + texture)."""

@classmethod

def define_schema(cls):

return io.Schema(

node_id="MoGePointMapToMesh",

search_aliases=["moge", "mesh", "geometry", "point map"],

display_name="Convert MoGe Point Map to Mesh",

description="Convert a MoGe point map into a 3D mesh.",

category="image/geometry estimation",

inputs=[

MoGeGeometry.Input("moge_geometry"),

io.Int.Input("batch_index", default=0, min=0, max=4096,

tooltip="Which image of a batched MoGe geometry to mesh. Per-image vertex counts "

"differ, so batches can't be stacked into a single MESH."),

io.Int.Input("decimation", default=1, min=1, max=8,

tooltip="Vertex stride; 1 = full resolution."),

io.Float.Input("discontinuity_threshold", default=0.04, min=0.0, max=1.0, step=0.01,

tooltip="Drop pixels whose 3x3 depth span exceeds this fraction. 0 = off."),

io.Boolean.Input("texture", default=True,

tooltip="Carry the source image through as the baseColor texture."),

],

outputs=[io.Mesh.Output()],

)

@classmethod

def execute(cls, moge_geometry, batch_index, decimation, discontinuity_threshold, texture) -> io.NodeOutput:

if "points" not in moge_geometry:

raise ValueError("moge_geometry has no points output.")

points = moge_geometry["points"]

B = points.shape[0]

if batch_index >= B:

raise ValueError(f"batch_index {batch_index} out of range; moge_geometry has batch size {B}.")

# Pass depth so the rtol edge check sees radial depth -- for panoramas

# points[..., 2] = cos(phi)*r goes negative below the equator and the rtol clamp would drop the bottom half.

edge_depth = moge_geometry["depth"][batch_index] if "depth" in moge_geometry else None

verts, faces, uvs = triangulate_grid_mesh(

points[batch_index], decimation=decimation,

discontinuity_threshold=discontinuity_threshold, depth=edge_depth,

)

if verts.shape[0] == 0 or faces.shape[0] == 0:

raise ValueError("MoGe produced an empty mesh; try discontinuity_threshold=0 or apply_mask=False.")

if "intrinsics" not in moge_geometry:

# Panorama: rotate MoGe spherical (Z up) -> glTF (Y up, Z back), correct for inside-the-sphere viewing)

verts = verts[:, [1, 2, 0]].contiguous()

else:

# Perspective MoGe (X right, Y down, Z forward) -> glTF; face flip keeps winding CCW after the Y/Z flip.

verts = verts * torch.tensor([1.0, -1.0, -1.0], dtype=verts.dtype)

faces = faces[:, [0, 2, 1]].contiguous()

tex = moge_geometry["image"][batch_index:batch_index + 1] if texture else None

mesh = Types.MESH(

vertices=verts.unsqueeze(0),

faces=faces.unsqueeze(0),

uvs=uvs.unsqueeze(0),

texture=tex,

)

"""Extract horizontal/vertical FOV from MoGe intrinsics, e.g. fov_y to feed SAM3DBody_Predict."""

@classmethod

def define_schema(cls):

return io.Schema(

node_id="MoGeGeometryToFOV",

search_aliases=["moge", "fov", "geometry", "intrinsics", "field of view"],

display_name="Get FoV from MoGe Geometry",

description="Derive the field of view and focal length from MoGe intrinsics.",

category="image/geometry estimation",

inputs=[

MoGeGeometry.Input("moge_geometry"),

io.Combo.Input("axis", options=["vertical", "horizontal", "diagonal"], default="vertical",

tooltip="'vertical' (fov_y), 'horizontal' (fov_x), or 'diagonal'."),

io.Combo.Input("unit", options=["degrees", "radians"], default="degrees",

tooltip="Output unit for the FOV."),

],

outputs=[

io.Float.Output(display_name="fov"),

io.Float.Output(display_name="focal_pixels"),

],

)

@classmethod

def execute(cls, moge_geometry, axis, unit) -> io.NodeOutput:

K = moge_geometry.get("intrinsics") if isinstance(moge_geometry, dict) else None

if K is None:

raise ValueError("moge_geometry has no intrinsics (panorama geometry has none).")

if K.ndim == 3:

K = K[0]

# MoGe normalizes fx by width and fy by height; with cx=cy=0.5 the half-extent

# in normalized units is 0.5, so fov = 2*atan(0.5 / f) per axis (hypot for diagonal).

hx = 0.5 / float(K[0, 0].item())

hy = 0.5 / float(K[1, 1].item())

half_tan = {"horizontal": hx, "vertical": hy, "diagonal": math.hypot(hx, hy)}[axis]

fov_radians = 2.0 * math.atan(half_tan)

fov = fov_radians if unit == "radians" else math.degrees(fov_radians)

# Pixels are square here, so fy*H == fx*W is the single lens focal in pixels.

src = next((moge_geometry[k] for k in ("image", "points", "depth") if k in moge_geometry), None)

if src is None:

raise ValueError("moge_geometry has no image/points/depth to read the pixel height from.")

H = int(src.shape[1])

focal_pixels = float(K[1, 1].item()) * H