Fix MoGe antialiased interpolation for half precision

This commit is contained in:
holys519 2026-06-26 03:25:04 +09:00
parent 639c8fa788
commit ee847d4ae0
2 changed files with 10 additions and 4 deletions

View File

@ -19,7 +19,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, _interpolate_antialias_safe, _view_plane_uv_grid
def _remap_points(points: torch.Tensor) -> torch.Tensor:
@ -68,9 +68,9 @@ class MoGeModelV1(nn.Module):
H, W = image.shape[-2:]
resize = ((num_tokens * 14 ** 2) / (H * W)) ** 0.5
rh, rw = int(H * resize), int(W * resize)
x = F.interpolate(image, (rh, rw), mode="bicubic", align_corners=False, antialias=True)
x = _interpolate_antialias_safe(image, (rh, rw), mode="bicubic", align_corners=False, antialias=True)
x = (x - self.image_mean) / self.image_std
x14 = F.interpolate(x, (rh // 14 * 14, rw // 14 * 14), mode="bilinear", align_corners=False, antialias=True)
x14 = _interpolate_antialias_safe(x, (rh // 14 * 14, rw // 14 * 14), mode="bilinear", align_corners=False, antialias=True)
n_layers = len(self.backbone.encoder.layer)
indices = list(range(n_layers - self.intermediate_layers, n_layers))

View File

@ -29,6 +29,12 @@ def _concat_view_plane_uv(x: torch.Tensor, aspect_ratio: float) -> torch.Tensor:
return torch.cat([x, uv], dim=1)
def _interpolate_antialias_safe(input: torch.Tensor, *args, **kwargs) -> torch.Tensor:
if kwargs.get("antialias") and input.dtype in (torch.float16, torch.bfloat16):
return F.interpolate(input.float(), *args, **kwargs).to(dtype=input.dtype)
return F.interpolate(input, *args, **kwargs)
class ResidualConvBlock(nn.Module):
def __init__(self, channels: int, hidden_channels: Optional[int] = None, in_norm: str = "layer_norm", hidden_norm: str = "group_norm",
dtype=None, device=None, operations=comfy.ops.manual_cast):
@ -135,7 +141,7 @@ class DINOv2Encoder(nn.Module):
def forward(self, image: torch.Tensor, token_rows: int, token_cols: int,
return_class_token: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
image_14 = F.interpolate(image, (token_rows * 14, token_cols * 14), mode="bilinear", align_corners=False, antialias=True)
image_14 = _interpolate_antialias_safe(image, (token_rows * 14, token_cols * 14), mode="bilinear", align_corners=False, antialias=True)
image_14 = (image_14 - self.image_mean) / self.image_std
feats = self.backbone.get_intermediate_layers(image_14, self.intermediate_layers, apply_norm=True)
x = torch.stack([