From a2f455c9da47fad15a09459e2a2a2de9ea50e7c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?= <40791699+kijai@users.noreply.github.com> Date: Thu, 17 Sep 2026 21:17:04 +0300 Subject: [PATCH] feat: Support MoGe 3 (CORE-443) (#16381) --- comfy/ldm/moge/model.py | 117 ++++++++++++++++++++++++----- comfy/ldm/moge/modules.py | 130 ++++++++++++++++++++++++++++++++- comfy/ldm/trellis2/flexgemm.py | 46 ++++++++++++ comfy_extras/nodes_moge.py | 14 +++- 4 files changed, 283 insertions(+), 24 deletions(-) diff --git a/comfy/ldm/moge/model.py b/comfy/ldm/moge/model.py index 8dd280c63..1b2c0ef9b 100644 --- a/comfy/ldm/moge/model.py +++ b/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() diff --git a/comfy/ldm/moge/modules.py b/comfy/ldm/moge/modules.py index 366411120..e9a622528 100644 --- a/comfy/ldm/moge/modules.py +++ b/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) diff --git a/comfy/ldm/trellis2/flexgemm.py b/comfy/ldm/trellis2/flexgemm.py index 76023992e..485c272c7 100644 --- a/comfy/ldm/trellis2/flexgemm.py +++ b/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] diff --git a/comfy_extras/nodes_moge.py b/comfy_extras/nodes_moge.py index 819421534..1861ef3ef 100644 --- a/comfy_extras/nodes_moge.py +++ b/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])