feat: Support MoGe 3 (CORE-443) (#16381)

comfy/ldm/moge/model.py

@@ -2,6 +2,7 @@
 
 V1: DINOv2 backbone + multi-output head (points, mask).
 V2: DINOv2 encoder + neck + per-output heads (points, mask, normal, optional metric-scale MLP).
+V3: V2 plus a sparse 3D UNet that iteratively refines the predicted log-depth.
 """
 
 
@@ -19,7 +20,7 @@ import comfy.model_patcher
 from comfy.image_encoders.dino2 import Dinov2Model
 
 from .geometry import depth_map_to_point_map, intrinsics_from_focal_center, recover_focal_shift
-from .modules import ConvStack, DINOv2Encoder, HeadV1, MLP, _view_plane_uv_grid
+from .modules import ConvStack, DINOv2Encoder, HeadV1, MLP, Sparse3DUNet, _view_plane_uv_grid
 
 
 def _remap_points(points: torch.Tensor) -> torch.Tensor:
@@ -30,7 +31,7 @@ def _remap_points(points: torch.Tensor) -> torch.Tensor:
 
 
 def _detect_dinov2(sd: dict, prefix: str) -> Dict[str, Any]:
-    # All shipped MoGe checkpoints use plain DINOv2
+    # All shipped MoGe checkpoints use plain DINOv2. ViT-g (MoGe-3) swaps the MLP for a fused SwiGLU.
     hidden = sd[prefix + "embeddings.cls_token"].shape[-1]
     layer_prefix = prefix + "encoder.layer."
     depth = 1 + max(int(k[len(layer_prefix):].split(".")[0]) for k in sd if k.startswith(layer_prefix))
@@ -39,7 +40,7 @@ def _detect_dinov2(sd: dict, prefix: str) -> Dict[str, Any]:
         "num_attention_heads": hidden // 64,
         "num_hidden_layers": depth,
         "layer_norm_eps": 1e-6,
-        "use_swiglu_ffn": False,
+        "use_swiglu_ffn": layer_prefix + "0.mlp.weights_in.weight" in sd,
     }
 
 
@@ -123,7 +124,8 @@ class MoGeModelV2(nn.Module):
         if normal_head is not None:
             self.normal_head = ConvStack(**normal_head, dtype=dtype, device=device, operations=operations)
 
-    def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]:
+    def _trunk(self, image: torch.Tensor, num_tokens: int) -> Tuple[List[torch.Tensor], torch.Tensor, torch.Tensor]:
+        """Encoder + neck. Returns (neck features, level-0 feature map, class token)."""
         B, _, H, W = image.shape
         device, dtype = image.device, image.dtype
         aspect_ratio = W / H
@@ -137,12 +139,15 @@ class MoGeModelV2(nn.Module):
                                     for L in range(5)]
         levels[0] = torch.cat([feat_top, levels[0]], dim=1)
 
-        feats = self.neck(levels)
+        return self.neck(levels), levels[0], cls_token
 
+    def _heads(self, feats: List[torch.Tensor], cls_token: torch.Tensor, raw_coord: torch.Tensor,
+               size: Tuple[int, int]) -> Dict[str, torch.Tensor]:
+        """Resize the head outputs to the image size and remap them into the public dict."""
         def _resize(v):
-            return F.interpolate(v, (H, W), mode="bilinear", align_corners=False)
+            return F.interpolate(v, size, mode="bilinear", align_corners=False)
 
-        points = _remap_points(_resize(self.points_head(feats)[-1]).permute(0, 2, 3, 1))
+        points = _remap_points(_resize(raw_coord).permute(0, 2, 3, 1))
         mask = _resize(self.mask_head(feats)[-1]).squeeze(1).sigmoid()
         metric_scale = self.scale_head(cls_token).squeeze(1).exp()
 
@@ -152,9 +157,20 @@ class MoGeModelV2(nn.Module):
             result["normal"] = F.normalize(normal.permute(0, 2, 3, 1), dim=-1)
         return result
 
+    def forward(self, image: torch.Tensor, num_tokens: int) -> Dict[str, torch.Tensor]:
+        feats, _conditioning, cls_token = self._trunk(image, num_tokens)
+        return self._heads(feats, cls_token, self.points_head(feats)[-1], image.shape[-2:])
+
     @classmethod
     def from_state_dict(cls, sd, dtype=None, device=None, operations=comfy.ops.manual_cast):
-        """Detect the v2 encoder/neck/heads config from sd, build a model, and load weights."""
+        """Detect the config from sd, build a model, and load weights."""
+        model = cls(**cls._detect_config(sd), dtype=dtype, device=device, operations=operations)
+        model.load_state_dict(sd, strict=True)
+        return model
+
+    @classmethod
+    def _detect_config(cls, sd) -> Dict[str, Any]:
+        """Reconstruct the v2 encoder/neck/heads config from the checkpoint keys."""
         backbone = _detect_dinov2(sd, prefix="encoder.backbone.")
         depth = backbone["num_hidden_layers"]
         n = cls.intermediate_layers
@@ -175,9 +191,7 @@ class MoGeModelV2(nn.Module):
         }
         if any(k.startswith("normal_head.") for k in sd):
             cfg["normal_head"] = cls._detect_convstack(sd, "normal_head.")
-        model = cls(**cfg, dtype=dtype, device=device, operations=operations)
-        model.load_state_dict(sd, strict=True)
-        return model
+        return cfg
 
     @staticmethod
     def _detect_convstack(sd: dict, prefix: str) -> Dict[str, Any]:
@@ -205,6 +219,56 @@ class MoGeModelV2(nn.Module):
         }
 
 
+class MoGeModelV3(MoGeModelV2):
+    """MoGe v3: the v2 architecture plus a sparse 3D UNet that iteratively refines the point map's log-depth."""
+
+    # Log-depth is binned at 1/256 to build the sparse volume, so the voxel grid stays finer
+    # than the depth detail the refiner is meant to recover.
+    refiner_depth_resolution = 256
+
+    def __init__(self, refiner: Dict[str, Any], dtype=None, device=None, operations=comfy.ops.manual_cast, **v2_kwargs):
+        super().__init__(**v2_kwargs, dtype=dtype, device=device, operations=operations)
+        self.refiner = Sparse3DUNet(**refiner, dtype=dtype, device=device, operations=operations)
+
+    def _refine_logz(self, coord: torch.Tensor, conditioning: torch.Tensor) -> torch.Tensor:
+        """One refinement pass over the point map at (x/z, y/z, logz), returning the updated logz."""
+        B, H, W, _ = coord.shape
+        device = coord.device
+        logz = coord[..., 2]
+
+        # Bin in fp32: logz * 256 lands where fp16's ULP exceeds 1 for far geometry, which would
+        # collapse neighbouring voxels. The refiner itself runs at the activation dtype.
+        z_bin = torch.round(logz.float() * self.refiner_depth_resolution).long()
+        z_bin = z_bin - z_bin.amin(dim=(1, 2), keepdim=True)
+
+        rows = torch.arange(H, device=device).view(1, H, 1).expand(B, H, W)
+        cols = torch.arange(W, device=device).view(1, 1, W).expand(B, H, W)
+        batch = torch.arange(B, device=device).view(B, 1, 1).expand(B, H, W)
+        coords = torch.stack([batch, rows, cols, z_bin], dim=-1).reshape(-1, 4).to(torch.int32)
+        spatial = (H, W, int(z_bin.amax()) + 1)
+
+        residual = self.refiner(coord.reshape(-1, 3), coords, spatial, conditioning)
+        return logz + residual.reshape(B, H, W)
+
+    def forward(self, image: torch.Tensor, num_tokens: int, refine_steps: int = 3) -> Dict[str, torch.Tensor]:
+        feats, conditioning, cls_token = self._trunk(image, num_tokens)
+
+        coord = self.points_head(feats)[-1].permute(0, 2, 3, 1)
+        for _ in range(refine_steps):
+            coord = torch.cat([coord[..., :2], self._refine_logz(coord, conditioning).unsqueeze(-1)], dim=-1)
+
+        # _remap_points takes exp(logz), which overflows fp16 past ~11.1, so hand the heads fp32.
+        return self._heads(feats, cls_token, coord.permute(0, 3, 1, 2).float(), image.shape[-2:])
+
+    @classmethod
+    def _detect_config(cls, sd) -> Dict[str, Any]:
+        cfg = super()._detect_config(sd)
+        # Both released MoGe-3 checkpoints share one refiner shape; only the conditioning
+        # width follows the encoder (1026 for ViT-L, 1538 for ViT-g).
+        cfg["refiner"] = {"encoder_channels": sd["refiner.encoder_fuse.weight"].shape[1]}
+        return cfg
+
+
 # Translate the Meta-style DINOv2 keys MoGe ships to the naming ComfyUI DINOv2 port expects,
 # and split each fused qkv tensor into Q/K/V.
 _DINOV2_TOPLEVEL_RENAMES = {
@@ -258,9 +322,14 @@ def _remap_state_dict(sd: dict) -> dict:
 
 
 def build_from_state_dict(sd: dict, dtype=None, device=None, operations=comfy.ops.manual_cast) -> nn.Module:
-    """Dispatch to v1 or v2 based on the DINOv2 backbone prefix."""
+    """Dispatch to v1, v2 or v3 based on the DINOv2 backbone prefix and the presence of the v3 refiner."""
     sd = _remap_state_dict(sd)
-    cls = MoGeModelV2 if any(k.startswith("encoder.backbone.") for k in sd) else MoGeModelV1
+    if not any(k.startswith("encoder.backbone.") for k in sd):
+        cls = MoGeModelV1
+    elif any(k.startswith("refiner.") for k in sd):
+        cls = MoGeModelV3
+    else:
+        cls = MoGeModelV2
     return cls.from_state_dict(sd, dtype=dtype, device=device, operations=operations)
 
 
@@ -274,7 +343,10 @@ class MoGeModel:
 
         self.model = build_from_state_dict(state_dict, dtype=self.dtype, device=offload_device, operations=comfy.ops.manual_cast).eval()
         self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
-        self.version = "v2" if hasattr(self.model, "encoder") else "v1"
+        if not hasattr(self.model, "encoder"):
+            self.version = "v1"
+        else:
+            self.version = "v3" if hasattr(self.model, "refiner") else "v2"
         self.mask_threshold = float(getattr(self.model, "mask_threshold", 0.5))
         nt = getattr(self.model, "num_tokens_range", (1200, 2500 if self.version == "v1" else 3600))
         self.num_tokens_range = (int(nt[0]), int(nt[1]))
@@ -282,11 +354,15 @@ class MoGeModel:
     def infer(self, image: torch.Tensor, num_tokens: Optional[int] = None,
               resolution_level: int = 9, fov_x: Optional[Union[Number, torch.Tensor]] = None,
               force_projection: bool = True, apply_mask: bool = True,
-              apply_metric_scale: bool = True
+              apply_metric_scale: bool = True, refine_steps: int = 3
               ) -> Dict[str, torch.Tensor]:
         """Run a single MoGe forward + post-process pass. image is (B, 3, H, W) in [0, 1]."""
         comfy.model_management.load_model_gpu(self.patcher)
-        image = image.to(device=self.load_device, dtype=torch.float32)
+
+        # Compute is fp32 or fp16 only: bf16 would cost 4x the error at the same speed
+        compute_dtype = self.dtype if self.dtype in (torch.float32, torch.float16) else torch.float16
+        activation_dtype = compute_dtype if self.version == "v3" else torch.float32
+        image = image.to(device=self.load_device, dtype=activation_dtype)
         H, W = image.shape[-2:]
         aspect_ratio = W / H
 
@@ -294,10 +370,13 @@ class MoGeModel:
             lo, hi = self.num_tokens_range
             num_tokens = int(lo + (resolution_level / 9) * (hi - lo))
 
-        out = self.model.forward(image, num_tokens=num_tokens)
+        # refine_steps only exists on v3; v1/v2 have no refiner to run.
+        extra = {"refine_steps": refine_steps} if self.version == "v3" else {}
+        out = self.model.forward(image, num_tokens=num_tokens, **extra)
         points = out["points"].float()  # recover_focal_shift goes through scipy on CPU; needs fp32.
         mask_binary = out["mask"] > self.mask_threshold
         normal = out.get("normal")
+        normal = normal.float() if normal is not None else None
         metric_scale = out.get("metric_scale")
 
         diag = (1 + aspect_ratio ** 2) ** 0.5
@@ -321,8 +400,8 @@ class MoGeModel:
         half = torch.tensor(0.5, device=points.device, dtype=points.dtype)
         intrinsics = intrinsics_from_focal_center(f_diag / aspect_ratio, f_diag, half, half)
         points[..., 2] = points[..., 2] + shift[..., None, None]
-        # v2 only: filter mask by depth>0 to drop metric-scale negative-depth artifacts.
-        if self.version == "v2":
+        # v2/v3 only: filter mask by depth>0 to drop metric-scale negative-depth artifacts.
+        if self.version != "v1":
             mask_binary = mask_binary & (points[..., 2] > 0)
         depth = points[..., 2].clone()
 

comfy/ldm/moge/modules.py

@@ -1,4 +1,4 @@
-"""Building blocks for MoGe: residual conv stack, resamplers, MLP, DINOv2 encoder, v1 head."""
+"""Building blocks for MoGe: residual conv stack, resamplers, MLP, DINOv2 encoder, v1 head, v3 sparse refiner."""
 
 
 from typing import List, Optional, Sequence, Tuple, Union
@@ -9,6 +9,7 @@ import torch.nn.functional as F
 
 import comfy.ops
 from comfy.image_encoders.dino2 import Dinov2Model
+from comfy.ldm.trellis2.flexgemm import sparse_pool3d_mean, sparse_submanifold_conv3d, sparse_upsample3d_nearest
 
 from .geometry import normalized_view_plane_uv
 
@@ -201,3 +202,130 @@ class HeadV1(nn.Module):
         x = F.interpolate(x, (img_h, img_w), mode="bilinear", align_corners=False)
         x = _concat_view_plane_uv(x, aspect)
         return [block(x) for block in self.output_block]
+
+
+class SubmanifoldConv3d(nn.Module):
+    """3x3x3 submanifold sparse conv. Weight is stored (C_out, K, K, K, C_in), as FlexGEMM writes it.
+
+    Kernel spatial axis i indexes coords column i + 1, matching FlexGEMM's neighbor map.
+    """
+
+    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3, dtype=None, device=None):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(out_channels, kernel_size, kernel_size, kernel_size, in_channels, dtype=dtype, device=device))
+        self.bias = nn.Parameter(torch.empty(out_channels, dtype=dtype, device=device))
+
+    def forward(self, feats, coords, spatial, neighbor_cache=None):
+        weight = comfy.ops.cast_to(self.weight, feats.dtype, feats.device)
+        bias = comfy.ops.cast_to(self.bias, feats.dtype, feats.device)
+        return sparse_submanifold_conv3d(feats, coords, spatial, weight, bias, neighbor_cache, (1, 1, 1))
+
+
+class SparseResBlock3d(nn.Module):
+    def __init__(self, channels: int, dtype=None, device=None, operations=comfy.ops.manual_cast):
+        super().__init__()
+        self.norm1 = operations.LayerNorm(channels, eps=1e-6, dtype=dtype, device=device)
+        self.conv1 = SubmanifoldConv3d(channels, channels, dtype=dtype, device=device)
+        self.conv2 = SubmanifoldConv3d(channels, channels, dtype=dtype, device=device)
+
+    def forward(self, feats, coords, spatial, neighbor_cache=None):
+        h = F.silu(self.norm1(feats))
+        h, neighbor_cache = self.conv1(h, coords, spatial, neighbor_cache)
+        h = F.silu(h)
+        h, neighbor_cache = self.conv2(h, coords, spatial, neighbor_cache)
+        return h + feats, neighbor_cache
+
+
+class PoolDown(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, factor: int, dtype=None, device=None, operations=comfy.ops.manual_cast):
+        super().__init__()
+        self.factor = factor
+        self.linear = operations.Linear(in_channels, out_channels, dtype=dtype, device=device)
+
+    def forward(self, feats, coords, spatial):
+        feats, coords, spatial, pool_index = sparse_pool3d_mean(feats, coords, spatial, self.factor)
+        return self.linear(feats), coords, spatial, pool_index
+
+
+class NearestUp(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, dtype=None, device=None, operations=comfy.ops.manual_cast):
+        super().__init__()
+        self.linear = operations.Linear(in_channels, out_channels, dtype=dtype, device=device)
+
+    def forward(self, feats, pool_index):
+        return sparse_upsample3d_nearest(self.linear(feats), pool_index)
+
+
+class Sparse3DUNet(nn.Module):
+    """MoGe v3 refiner: sparse 3D UNet over the voxelized point map, conditioned on the ViT feature map.
+
+    Takes the sparse volume as (feats, coords, spatial) where coords are (batch, row, col, z_bin),
+    and returns one residual per input voxel.
+    """
+
+    def __init__(self, encoder_channels: int, in_channels: int = 3, out_channels: int = 1,
+                 model_channels: Sequence[int] = (32, 64, 128, 256, 512), blocks_per_level: int = 1,
+                 factor: int = 2, dtype=None, device=None, operations=comfy.ops.manual_cast):
+        super().__init__()
+        self.factor = factor
+        kwargs = {"dtype": dtype, "device": device, "operations": operations}
+        # (shallow, deep) channel pair per resolution transition
+        pairs = list(zip(model_channels[:-1], model_channels[1:]))
+
+        def stage(channels):
+            return nn.ModuleList([SparseResBlock3d(channels, **kwargs) for _ in range(blocks_per_level)])
+
+        self.input_proj = operations.Linear(in_channels, model_channels[0], dtype=dtype, device=device)
+        self.encoder_fuse = operations.Linear(encoder_channels, model_channels[-1], dtype=dtype, device=device)
+        self.fuse_proj = nn.Sequential(
+            operations.Linear(model_channels[-1] * 2, model_channels[-1], dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Linear(model_channels[-1], model_channels[-1], dtype=dtype, device=device),
+        )
+
+        self.down_stages = nn.ModuleList([stage(ch) for ch in model_channels])
+        self.downsample_blocks = nn.ModuleList([PoolDown(lo, hi, factor, **kwargs) for lo, hi in pairs])
+        self.bottleneck_stage = stage(model_channels[-1])
+        # The decoder runs deepest-first, so it walks the transitions in reverse.
+        self.upsample_blocks = nn.ModuleList([NearestUp(hi, lo, **kwargs) for lo, hi in reversed(pairs)])
+        self.up_stages = nn.ModuleList([stage(lo) for lo, _ in reversed(pairs)])
+        self.out_proj = operations.Linear(model_channels[0], out_channels, dtype=dtype, device=device)
+
+    def forward(self, feats, coords, spatial, encoder_feature):
+        num_levels = len(self.down_stages)
+        num_transitions = len(self.downsample_blocks)
+        # Coords at level k are identical on the down and up passes, so the submanifold
+        # neighbor map built on the way down is still valid on the way back up.
+        conv_caches: List[Optional[torch.Tensor]] = [None] * num_levels
+        pool_indices: List[Optional[torch.Tensor]] = [None] * num_transitions
+        skips: List[Optional[Tuple[torch.Tensor, torch.Tensor, tuple]]] = [None] * num_transitions
+
+        feats = self.input_proj(feats)
+
+        for i, blocks in enumerate(self.down_stages):
+            cache = conv_caches[i]
+            for block in blocks:
+                feats, cache = block(feats, coords, spatial, cache)
+            conv_caches[i] = cache
+            if i < num_transitions:
+                skips[i] = (feats, coords, spatial)
+                feats, coords, spatial, pool_indices[i] = self.downsample_blocks[i](feats, coords, spatial)
+
+        conditioning = encoder_feature[coords[:, 0].long(), :, coords[:, 1].long(), coords[:, 2].long()]
+        feats = self.fuse_proj(torch.cat([feats, self.encoder_fuse(conditioning)], dim=-1))
+
+        cache = conv_caches[num_levels - 1]
+        for block in self.bottleneck_stage:
+            feats, cache = block(feats, coords, spatial, cache)
+        conv_caches[num_levels - 1] = cache
+
+        for i, (upsample, blocks) in enumerate(zip(self.upsample_blocks, self.up_stages)):
+            level = num_levels - 2 - i
+            skip_feats, coords, spatial = skips[level]
+            feats = upsample(feats, pool_indices[level]) + skip_feats
+            cache = conv_caches[level]
+            for block in blocks:
+                feats, cache = block(feats, coords, spatial, cache)
+            conv_caches[level] = cache
+
+        return self.out_proj(feats)

comfy/ldm/trellis2/flexgemm.py

@@ -164,3 +164,49 @@ def sparse_submanifold_conv3d(
         output += bias.unsqueeze(0).to(output.dtype)
 
     return output, neighbor
+
+
+def sparse_pool3d_mean(
+    feats: torch.Tensor,
+    coords: torch.Tensor,
+    shape: tuple,
+    factor: int,
+) -> Tuple[torch.Tensor, torch.Tensor, tuple, torch.Tensor]:
+    """Average-pool a sparse volume by `factor` along the three spatial axes.
+
+    coords are (batch, x, y, z) and shape is the spatial extent (X, Y, Z), as for
+    sparse_submanifold_conv3d. Returns the pooled (feats, coords, shape) plus the
+    fine->coarse index, which is what sparse_upsample3d_nearest needs to invert the pooling.
+    """
+    batch = coords[:, 0].long()
+    fine = coords[:, 1:4].long()
+    out_x, out_y, out_z = ((int(s) + factor - 1) // factor for s in shape)
+
+    coarse = fine.div(factor, rounding_mode="floor")
+    flat = ((batch * out_x + coarse[:, 0]) * out_y + coarse[:, 1]) * out_z + coarse[:, 2]
+    unique_flat, index = torch.unique(flat, return_inverse=True)
+
+    pooled = torch.zeros((unique_flat.shape[0], feats.shape[-1]), device=feats.device, dtype=feats.dtype)
+    pooled.index_add_(0, index, feats)
+    counts = torch.zeros((unique_flat.shape[0],), device=feats.device, dtype=feats.dtype)
+    counts.index_add_(0, index, torch.ones_like(index, dtype=feats.dtype))
+    pooled /= counts.unsqueeze(-1)
+
+    z = unique_flat % out_z
+    rest = unique_flat.div(out_z, rounding_mode="floor")
+    y = rest % out_y
+    rest = rest.div(out_y, rounding_mode="floor")
+    x = rest % out_x
+    b = rest.div(out_x, rounding_mode="floor")
+    out_coords = torch.stack([b, x, y, z], dim=-1).to(torch.int32)
+
+    return pooled, out_coords, (out_x, out_y, out_z), index
+
+
+def sparse_upsample3d_nearest(feats: torch.Tensor, pool_index: torch.Tensor) -> torch.Tensor:
+    """Invert sparse_pool3d_mean's spatial mapping: each fine voxel takes its parent's feature.
+
+    `pool_index` is the index returned by the matching sparse_pool3d_mean call, so the
+    output is aligned with that call's input coords (the UNet's skip coords).
+    """
+    return feats[pool_index]

comfy_extras/nodes_moge.py

@@ -106,12 +106,14 @@ class MoGePanoramaInference(io.ComfyNode):
                              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) -> io.NodeOutput:
+    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]})")
@@ -155,7 +157,8 @@ class MoGePanoramaInference(io.ComfyNode):
                 # 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)
+                                          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]
@@ -228,12 +231,14 @@ class MoGeInference(io.ComfyNode):
                 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) -> io.NodeOutput:
+    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()
@@ -246,7 +251,8 @@ class MoGeInference(io.ComfyNode):
             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))
+                                               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])