mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-06 22:51:18 +08:00
Fix MoGe antialiased interpolation for half precision
This commit is contained in:
parent
639c8fa788
commit
ee847d4ae0
@ -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))
|
||||
|
||||
@ -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([
|
||||
|
||||
Loading…
Reference in New Issue
Block a user