LUCAS M MORGAN
dfd7e68cbb
Merge LoadSharpModel and SharpPredict into this pack; register all nodes from __init__.py. Bundle web/js/gsplat-bundle.js and resolve viewer load via import.meta.url. Add LICENSE, THIRD_PARTY_NOTICES, and GPL-3.0 text. Update README and requirements (huggingface-hub). Remove GeometryPack runtime dependency for the WebGL viewer. Made-with: Cursor
1985 lines
72 KiB
Python
1985 lines
72 KiB
Python
"""All neural network modules for SHARP monocular 3D Gaussian Splatting predictor.
|
|
|
|
Consolidated from ~40 files into a single ComfyUI-native module.
|
|
All learnable layers use operations= for ComfyUI weight management.
|
|
Custom ViT replaces timm dependency with matching key names for checkpoint compatibility.
|
|
|
|
For licensing see accompanying LICENSE file.
|
|
Copyright (C) 2025 Apple Inc. All Rights Reserved.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import abc
|
|
import copy
|
|
import math
|
|
from typing import Literal, NamedTuple, Sequence, Tuple, Union
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
|
|
import comfy.model_management
|
|
import comfy.ops
|
|
from comfy.ldm.modules.attention import optimized_attention
|
|
|
|
from .color_space import ColorSpace, sRGB2linearRGB
|
|
from .gaussians import Gaussians3D
|
|
|
|
# Activation helpers (replaces math.py abstraction — only sigmoid/exp used at inference)
|
|
def _inverse_sigmoid(x: torch.Tensor) -> torch.Tensor:
|
|
return torch.log(x / (1 - x))
|
|
|
|
def _inverse_softplus(x: torch.Tensor) -> torch.Tensor:
|
|
return torch.log(torch.exp(x) - 1)
|
|
|
|
_ACTIVATIONS = {
|
|
"sigmoid": (torch.sigmoid, _inverse_sigmoid),
|
|
"exp": (torch.exp, torch.log),
|
|
}
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Type aliases
|
|
# ---------------------------------------------------------------------------
|
|
|
|
NormLayerName = Literal["noop", "batch_norm", "group_norm", "instance_norm"]
|
|
UpsamplingMode = Literal["transposed_conv", "nearest", "bilinear"]
|
|
DimsDecoder = Tuple[int, int, int, int, int]
|
|
DPTImageEncoderType = Literal["skip_conv", "skip_conv_kernel2"]
|
|
ColorInitOption = Literal["none", "first_layer", "all_layers"]
|
|
DepthInitOption = Literal["surface_min", "surface_max", "base_depth", "linear_disparity"]
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helper functions
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def norm_layer_2d(
|
|
num_features: int,
|
|
norm_type: NormLayerName,
|
|
num_groups: int = 8,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> nn.Module:
|
|
"""Create normalization layer."""
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if norm_type == "noop":
|
|
return nn.Identity()
|
|
elif norm_type == "batch_norm":
|
|
return nn.BatchNorm2d(num_features=num_features)
|
|
elif norm_type == "group_norm":
|
|
return operations.GroupNorm(
|
|
num_channels=num_features, num_groups=num_groups, dtype=dtype, device=device
|
|
)
|
|
elif norm_type == "instance_norm":
|
|
return nn.InstanceNorm2d(num_features=num_features)
|
|
else:
|
|
raise ValueError(f"Invalid normalization layer type: {norm_type}")
|
|
|
|
|
|
def upsampling_layer(
|
|
upsampling_mode: UpsamplingMode,
|
|
scale_factor: int,
|
|
dim_in: int,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> nn.Module:
|
|
"""Create upsampling layer."""
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if upsampling_mode == "transposed_conv":
|
|
return operations.ConvTranspose2d(
|
|
in_channels=dim_in,
|
|
out_channels=dim_in,
|
|
kernel_size=scale_factor,
|
|
stride=scale_factor,
|
|
padding=0,
|
|
bias=False,
|
|
dtype=dtype,
|
|
device=device,
|
|
)
|
|
elif upsampling_mode in ("nearest", "bilinear"):
|
|
return nn.Upsample(scale_factor=scale_factor, mode=upsampling_mode)
|
|
else:
|
|
raise ValueError(f"Invalid upsampling mode {upsampling_mode}.")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# NamedTuples
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class GaussianBaseValues(NamedTuple):
|
|
"""Base values for gaussian predictor."""
|
|
|
|
mean_x_ndc: torch.Tensor
|
|
mean_y_ndc: torch.Tensor
|
|
mean_inverse_z_ndc: torch.Tensor
|
|
scales: torch.Tensor
|
|
quaternions: torch.Tensor
|
|
colors: torch.Tensor
|
|
opacities: torch.Tensor
|
|
|
|
|
|
class InitializerOutput(NamedTuple):
|
|
"""Output of initializer."""
|
|
|
|
gaussian_base_values: GaussianBaseValues
|
|
feature_input: torch.Tensor
|
|
global_scale: torch.Tensor | None = None
|
|
|
|
|
|
class MonodepthOutput(NamedTuple):
|
|
"""Output of the monodepth model."""
|
|
|
|
disparity: torch.Tensor
|
|
encoder_features: list[torch.Tensor]
|
|
decoder_features: torch.Tensor
|
|
output_features: list[torch.Tensor]
|
|
intermediate_features: list[torch.Tensor] = []
|
|
|
|
|
|
class ImageFeatures(NamedTuple):
|
|
"""Image feature extracted from decoder."""
|
|
|
|
texture_features: torch.Tensor
|
|
geometry_features: torch.Tensor
|
|
|
|
|
|
# ===========================================================================
|
|
# Vision Transformer (replaces timm dependency)
|
|
# Attribute names match timm exactly for checkpoint key compatibility.
|
|
# ===========================================================================
|
|
|
|
|
|
class PatchEmbed(nn.Module):
|
|
"""2D Image to Patch Embedding. Matches timm PatchEmbed key naming."""
|
|
|
|
def __init__(
|
|
self,
|
|
img_size: int = 384,
|
|
patch_size: int = 16,
|
|
in_chans: int = 3,
|
|
embed_dim: int = 1024,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.img_size = (img_size, img_size) if isinstance(img_size, int) else img_size
|
|
self.patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size
|
|
self.grid_size = (
|
|
self.img_size[0] // self.patch_size[0],
|
|
self.img_size[1] // self.patch_size[1],
|
|
)
|
|
self.num_patches = self.grid_size[0] * self.grid_size[1]
|
|
# Key: patch_embed.proj.weight / .bias
|
|
self.proj = operations.Conv2d(
|
|
in_chans, embed_dim, kernel_size=patch_size, stride=patch_size,
|
|
dtype=dtype, device=device,
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = self.proj(x)
|
|
x = x.flatten(2).transpose(1, 2) # [B, C, H, W] -> [B, N, C]
|
|
return x
|
|
|
|
|
|
class LayerScale(nn.Module):
|
|
"""Per-channel learnable scaling. Key: ls{1,2}.gamma"""
|
|
|
|
def __init__(self, dim: int, init_values: float = 1e-5, dtype=None, device=None):
|
|
super().__init__()
|
|
self.gamma = nn.Parameter(
|
|
torch.full((dim,), init_values, dtype=dtype, device=device)
|
|
)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
return x * self.gamma
|
|
|
|
|
|
class ViTAttention(nn.Module):
|
|
"""ViT self-attention with fused qkv and ComfyUI optimized_attention.
|
|
|
|
Keys: attn.qkv.{weight,bias}, attn.proj.{weight,bias}
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
qkv_bias: bool = True,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.num_heads = num_heads
|
|
self.qkv = operations.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
|
|
self.proj = operations.Linear(dim, dim, dtype=dtype, device=device)
|
|
|
|
def forward(self, x: torch.Tensor, transformer_options={}) -> torch.Tensor:
|
|
qkv = self.qkv(x)
|
|
q, k, v = qkv.chunk(3, dim=-1) # Each [B, N, C]
|
|
out = optimized_attention(q, k, v, heads=self.num_heads,
|
|
transformer_options=transformer_options)
|
|
return self.proj(out)
|
|
|
|
|
|
class ViTMlp(nn.Module):
|
|
"""ViT MLP supporting vanilla (GELU) and GLU (GEGLU) modes.
|
|
|
|
Keys: mlp.fc1.{weight,bias}, mlp.fc2.{weight,bias}
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
mlp_ratio: float = 4.0,
|
|
use_glu: bool = False,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
hidden = int(dim * mlp_ratio)
|
|
self.use_glu = use_glu
|
|
if use_glu:
|
|
# GluMlp: fc1 outputs hidden (chunked into 2 halves for gating)
|
|
self.fc1 = operations.Linear(dim, hidden, dtype=dtype, device=device)
|
|
self.fc2 = operations.Linear(hidden // 2, dim, dtype=dtype, device=device)
|
|
else:
|
|
self.fc1 = operations.Linear(dim, hidden, dtype=dtype, device=device)
|
|
self.fc2 = operations.Linear(hidden, dim, dtype=dtype, device=device)
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
x = self.fc1(x)
|
|
if self.use_glu:
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
x = x1 * F.gelu(x2)
|
|
else:
|
|
x = F.gelu(x)
|
|
return self.fc2(x)
|
|
|
|
|
|
class VisionTransformerBlock(nn.Module):
|
|
"""Single ViT block. Keys match timm Block naming exactly."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim: int,
|
|
num_heads: int,
|
|
mlp_ratio: float = 4.0,
|
|
qkv_bias: bool = True,
|
|
init_values: float = 1e-5,
|
|
use_glu_mlp: bool = False,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.norm1 = operations.LayerNorm(dim, dtype=dtype, device=device)
|
|
self.attn = ViTAttention(
|
|
dim, num_heads, qkv_bias=qkv_bias,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.ls1 = LayerScale(dim, init_values=init_values, dtype=dtype, device=device)
|
|
self.norm2 = operations.LayerNorm(dim, dtype=dtype, device=device)
|
|
self.mlp = ViTMlp(
|
|
dim, mlp_ratio=mlp_ratio, use_glu=use_glu_mlp,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.ls2 = LayerScale(dim, init_values=init_values, dtype=dtype, device=device)
|
|
|
|
def forward(self, x: torch.Tensor, transformer_options={}) -> torch.Tensor:
|
|
x = x + self.ls1(self.attn(self.norm1(x), transformer_options=transformer_options))
|
|
x = x + self.ls2(self.mlp(self.norm2(x)))
|
|
return x
|
|
|
|
|
|
class VisionTransformer(nn.Module):
|
|
"""Custom ViT matching timm VisionTransformer state_dict keys exactly.
|
|
|
|
Replaces TimmViT(timm.models.VisionTransformer) with full operations= support.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
img_size: int = 384,
|
|
patch_size: int = 16,
|
|
in_chans: int = 3,
|
|
embed_dim: int = 1024,
|
|
depth: int = 24,
|
|
num_heads: int = 16,
|
|
mlp_ratio: float = 4.0,
|
|
qkv_bias: bool = True,
|
|
init_values: float = 1e-5,
|
|
use_glu_mlp: bool = False,
|
|
num_classes: int = 0,
|
|
intermediate_features_ids: list[int] | None = None,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.embed_dim = embed_dim
|
|
self.dim_in = in_chans
|
|
self.num_prefix_tokens = 1 # CLS token
|
|
self.intermediate_features_ids = intermediate_features_ids
|
|
|
|
# Patch embedding
|
|
self.patch_embed = PatchEmbed(
|
|
img_size=img_size, patch_size=patch_size, in_chans=in_chans,
|
|
embed_dim=embed_dim, dtype=dtype, device=device, operations=operations,
|
|
)
|
|
num_patches = self.patch_embed.num_patches
|
|
|
|
# CLS token and positional embedding (nn.Parameter, not operations)
|
|
self.cls_token = nn.Parameter(
|
|
torch.zeros(1, 1, embed_dim, dtype=dtype, device=device)
|
|
)
|
|
self.pos_embed = nn.Parameter(
|
|
torch.zeros(1, num_patches + 1, embed_dim, dtype=dtype, device=device)
|
|
)
|
|
|
|
# Pre-norm and patch drop (Identity for DINOv2)
|
|
self.norm_pre = nn.Identity()
|
|
self.patch_drop = nn.Identity()
|
|
|
|
# Transformer blocks
|
|
self.blocks = nn.ModuleList([
|
|
VisionTransformerBlock(
|
|
dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio,
|
|
qkv_bias=qkv_bias, init_values=init_values, use_glu_mlp=use_glu_mlp,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
for _ in range(depth)
|
|
])
|
|
|
|
# Final norm
|
|
self.norm = operations.LayerNorm(embed_dim, dtype=dtype, device=device)
|
|
|
|
# Classification head (unused at inference, but needed for checkpoint key compatibility)
|
|
if num_classes > 0:
|
|
self.head = operations.Linear(embed_dim, num_classes, dtype=dtype, device=device)
|
|
else:
|
|
self.head = nn.Identity()
|
|
|
|
def reshape_feature(self, embeddings: torch.Tensor) -> torch.Tensor:
|
|
"""Discard class token and reshape 1D feature map to a 2D grid."""
|
|
batch_size, seq_len, channel = embeddings.shape
|
|
height, width = self.patch_embed.grid_size
|
|
|
|
if self.num_prefix_tokens:
|
|
embeddings = embeddings[:, self.num_prefix_tokens:, :]
|
|
|
|
embeddings = embeddings.reshape(batch_size, height, width, channel).permute(0, 3, 1, 2)
|
|
return embeddings
|
|
|
|
def forward(
|
|
self, input_tensor: torch.Tensor, transformer_options={}
|
|
) -> tuple[torch.Tensor, dict[int, torch.Tensor]]:
|
|
"""Forward pass with intermediate feature extraction."""
|
|
intermediate_features: dict[int, torch.Tensor] = {}
|
|
|
|
x = self.patch_embed(input_tensor)
|
|
|
|
# Prepend CLS token and add positional embedding
|
|
cls_tokens = self.cls_token.expand(x.shape[0], -1, -1)
|
|
x = torch.cat((cls_tokens, x), dim=1)
|
|
x = x + self.pos_embed
|
|
|
|
x = self.patch_drop(x)
|
|
x = self.norm_pre(x)
|
|
|
|
for idx, block in enumerate(self.blocks):
|
|
x = block(x, transformer_options=transformer_options)
|
|
if (
|
|
self.intermediate_features_ids is not None
|
|
and idx in self.intermediate_features_ids
|
|
):
|
|
intermediate_features[idx] = x
|
|
|
|
x = self.norm(x)
|
|
x = self.reshape_feature(x)
|
|
return x, intermediate_features
|
|
|
|
def internal_resolution(self) -> int:
|
|
"""Return the internal image size of the network."""
|
|
if isinstance(self.patch_embed.img_size, tuple):
|
|
return self.patch_embed.img_size[0]
|
|
return self.patch_embed.img_size
|
|
|
|
|
|
# ===========================================================================
|
|
# Sliding Pyramid Network
|
|
# ===========================================================================
|
|
|
|
|
|
@torch.fx.wrap
|
|
def split(
|
|
image: torch.Tensor, overlap_ratio: float = 0.25, patch_size: int = 384
|
|
) -> torch.Tensor:
|
|
"""Split the input into small patches with sliding window."""
|
|
patch_stride = int(patch_size * (1 - overlap_ratio))
|
|
image_size = image.shape[-1]
|
|
steps = int(math.ceil((image_size - patch_size) / patch_stride)) + 1
|
|
|
|
x_patch_list = []
|
|
for j in range(steps):
|
|
j0 = j * patch_stride
|
|
j1 = j0 + patch_size
|
|
for i in range(steps):
|
|
i0 = i * patch_stride
|
|
i1 = i0 + patch_size
|
|
x_patch_list.append(image[..., j0:j1, i0:i1])
|
|
|
|
return torch.cat(x_patch_list, dim=0)
|
|
|
|
|
|
@torch.fx.wrap
|
|
def merge(
|
|
image_patches: torch.Tensor, batch_size: int, padding: int = 3
|
|
) -> torch.Tensor:
|
|
"""Merge the patched input into a image with sliding window."""
|
|
steps = int(math.sqrt(image_patches.shape[0] // batch_size))
|
|
idx = 0
|
|
|
|
output_list = []
|
|
for j in range(steps):
|
|
output_row_list = []
|
|
for i in range(steps):
|
|
output = image_patches[batch_size * idx : batch_size * (idx + 1)]
|
|
if padding != 0:
|
|
if j != 0:
|
|
output = output[..., padding:, :]
|
|
if i != 0:
|
|
output = output[..., :, padding:]
|
|
if j != steps - 1:
|
|
output = output[..., :-padding, :]
|
|
if i != steps - 1:
|
|
output = output[..., :, :-padding]
|
|
output_row_list.append(output)
|
|
idx += 1
|
|
output_row = torch.cat(output_row_list, dim=-1)
|
|
output_list.append(output_row)
|
|
return torch.cat(output_list, dim=-2)
|
|
|
|
|
|
def _create_spn_project_upsample_block(
|
|
dim_in: int,
|
|
dim_out: int,
|
|
upsample_layers: int,
|
|
dim_intermediate: int | None = None,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> nn.Module:
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if dim_intermediate is None:
|
|
dim_intermediate = dim_out
|
|
blocks: list[nn.Module] = [
|
|
operations.Conv2d(
|
|
in_channels=dim_in, out_channels=dim_intermediate,
|
|
kernel_size=1, stride=1, padding=0, bias=False,
|
|
dtype=dtype, device=device,
|
|
)
|
|
]
|
|
blocks += [
|
|
operations.ConvTranspose2d(
|
|
in_channels=dim_intermediate if i == 0 else dim_out,
|
|
out_channels=dim_out, kernel_size=2, stride=2, padding=0, bias=False,
|
|
dtype=dtype, device=device,
|
|
)
|
|
for i in range(upsample_layers)
|
|
]
|
|
return nn.Sequential(*blocks)
|
|
|
|
|
|
class SlidingPyramidNetwork(nn.Module):
|
|
"""Sliding Pyramid Network for multi-resolution ViT encodings."""
|
|
|
|
def __init__(
|
|
self,
|
|
dims_encoder: list[int],
|
|
patch_encoder: VisionTransformer,
|
|
image_encoder: VisionTransformer,
|
|
use_patch_overlap: bool = True,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.dim_in = patch_encoder.dim_in
|
|
self.dims_encoder = list(dims_encoder)
|
|
self.patch_encoder = patch_encoder
|
|
self.image_encoder = image_encoder
|
|
|
|
base_embed_dim = patch_encoder.embed_dim
|
|
lowres_embed_dim = image_encoder.embed_dim
|
|
self.patch_size = patch_encoder.internal_resolution()
|
|
self.use_patch_overlap = use_patch_overlap
|
|
|
|
self.patch_intermediate_features_ids = patch_encoder.intermediate_features_ids
|
|
if (
|
|
not isinstance(self.patch_intermediate_features_ids, list)
|
|
or not len(self.patch_intermediate_features_ids) == 4
|
|
):
|
|
raise ValueError("Patch intermediate feature ids must be a 4-item list.")
|
|
|
|
self.image_intermediate_features_ids = image_encoder.intermediate_features_ids
|
|
|
|
self.upsample_latent0 = _create_spn_project_upsample_block(
|
|
dim_in=base_embed_dim, dim_out=self.dims_encoder[0],
|
|
upsample_layers=3, dim_intermediate=self.dims_encoder[1],
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.upsample_latent1 = _create_spn_project_upsample_block(
|
|
dim_in=base_embed_dim, dim_out=self.dims_encoder[1], upsample_layers=2,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.upsample0 = _create_spn_project_upsample_block(
|
|
dim_in=base_embed_dim, dim_out=self.dims_encoder[2], upsample_layers=1,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.upsample1 = _create_spn_project_upsample_block(
|
|
dim_in=base_embed_dim, dim_out=self.dims_encoder[3], upsample_layers=1,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.upsample2 = _create_spn_project_upsample_block(
|
|
dim_in=base_embed_dim, dim_out=self.dims_encoder[4], upsample_layers=1,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
|
|
self.upsample_lowres = operations.ConvTranspose2d(
|
|
in_channels=lowres_embed_dim, out_channels=self.dims_encoder[4],
|
|
kernel_size=2, stride=2, padding=0, bias=True,
|
|
dtype=dtype, device=device,
|
|
)
|
|
self.fuse_lowres = operations.Conv2d(
|
|
in_channels=(self.dims_encoder[4] + self.dims_encoder[4]),
|
|
out_channels=self.dims_encoder[4],
|
|
kernel_size=1, stride=1, padding=0, bias=True,
|
|
dtype=dtype, device=device,
|
|
)
|
|
|
|
def internal_resolution(self) -> int:
|
|
return self.patch_size * 4
|
|
|
|
def _create_pyramid(
|
|
self, x: torch.Tensor
|
|
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
x0 = x
|
|
x1 = F.interpolate(x, size=None, scale_factor=0.5, mode="bilinear", align_corners=False)
|
|
x2 = F.interpolate(x, size=None, scale_factor=0.25, mode="bilinear", align_corners=False)
|
|
return x0, x1, x2
|
|
|
|
def forward(self, x: torch.Tensor) -> list[torch.Tensor]:
|
|
batch_size = x.shape[0]
|
|
x0, x1, x2 = self._create_pyramid(x)
|
|
|
|
if self.use_patch_overlap:
|
|
x0_patches = split(x0, overlap_ratio=0.25, patch_size=self.patch_size)
|
|
x1_patches = split(x1, overlap_ratio=0.5, patch_size=self.patch_size)
|
|
x2_patches = x2
|
|
padding = 3
|
|
else:
|
|
x0_patches = split(x0, overlap_ratio=0.0, patch_size=self.patch_size)
|
|
x1_patches = split(x1, overlap_ratio=0.0, patch_size=self.patch_size)
|
|
x2_patches = x2
|
|
padding = 0
|
|
x0_tile_size = x0_patches.shape[0]
|
|
|
|
x_pyramid_patches = torch.cat((x0_patches, x1_patches, x2_patches), dim=0)
|
|
|
|
# Dynamic chunking based on free VRAM (same pattern as ComfyUI's attention_split)
|
|
N = x_pyramid_patches.shape[0]
|
|
mem_free = comfy.model_management.get_free_memory(x_pyramid_patches.device)
|
|
patch_mem_estimate = 200 * 1024 * 1024 # ~200 MB per patch through ViT
|
|
chunk_size = max(1, int(mem_free * 0.8 / patch_mem_estimate))
|
|
chunk_size = min(chunk_size, N)
|
|
|
|
if chunk_size >= N:
|
|
x_pyramid_encodings, patch_intermediate_features = self.patch_encoder(x_pyramid_patches)
|
|
else:
|
|
encoding_chunks = []
|
|
intermediate_chunks: dict[int, list[torch.Tensor]] = {}
|
|
for i in range(0, N, chunk_size):
|
|
comfy.model_management.throw_exception_if_processing_interrupted()
|
|
enc, inter = self.patch_encoder(x_pyramid_patches[i:i + chunk_size])
|
|
encoding_chunks.append(enc)
|
|
for layer_id, feat in inter.items():
|
|
intermediate_chunks.setdefault(layer_id, []).append(feat)
|
|
x_pyramid_encodings = torch.cat(encoding_chunks, dim=0)
|
|
patch_intermediate_features = {
|
|
layer_id: torch.cat(feats, dim=0)
|
|
for layer_id, feats in intermediate_chunks.items()
|
|
}
|
|
|
|
# Merge intermediate features
|
|
x_latent0_encodings = self.patch_encoder.reshape_feature(
|
|
patch_intermediate_features[self.patch_intermediate_features_ids[0]]
|
|
)
|
|
x_latent0_features = merge(
|
|
x_latent0_encodings[: batch_size * x0_tile_size],
|
|
batch_size=batch_size, padding=padding,
|
|
)
|
|
|
|
x_latent1_encodings = self.patch_encoder.reshape_feature(
|
|
patch_intermediate_features[self.patch_intermediate_features_ids[1]]
|
|
)
|
|
x_latent1_features = merge(
|
|
x_latent1_encodings[: batch_size * x0_tile_size],
|
|
batch_size=batch_size, padding=padding,
|
|
)
|
|
|
|
x0_encodings, x1_encodings, x2_encodings = torch.split(
|
|
x_pyramid_encodings,
|
|
[len(x0_patches), len(x1_patches), len(x2_patches)],
|
|
dim=0,
|
|
)
|
|
|
|
x0_features = merge(x0_encodings, batch_size=batch_size, padding=padding)
|
|
x1_features = merge(x1_encodings, batch_size=batch_size, padding=2 * padding)
|
|
x2_features = x2_encodings
|
|
|
|
x_lowres_features, _ = self.image_encoder(x2_patches)
|
|
|
|
x_latent0_features = self.upsample_latent0(x_latent0_features)
|
|
x_latent1_features = self.upsample_latent1(x_latent1_features)
|
|
x0_features = self.upsample0(x0_features)
|
|
x1_features = self.upsample1(x1_features)
|
|
x2_features = self.upsample2(x2_features)
|
|
x_lowres_features = self.upsample_lowres(x_lowres_features)
|
|
x_lowres_features = self.fuse_lowres(
|
|
torch.cat((x2_features, x_lowres_features), dim=1)
|
|
)
|
|
|
|
return [
|
|
x_latent0_features,
|
|
x_latent1_features,
|
|
x0_features,
|
|
x1_features,
|
|
x_lowres_features,
|
|
]
|
|
|
|
|
|
# ===========================================================================
|
|
# Blocks (ResidualBlock, FeatureFusionBlock2d)
|
|
# ===========================================================================
|
|
|
|
|
|
class ResidualBlock(nn.Module):
|
|
"""Generic residual block."""
|
|
|
|
def __init__(self, residual: nn.Module, shortcut: nn.Module | None = None) -> None:
|
|
super().__init__()
|
|
self.residual = residual
|
|
self.shortcut = shortcut
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
delta_x = self.residual(x)
|
|
if self.shortcut is not None:
|
|
x = self.shortcut(x)
|
|
return x + delta_x
|
|
|
|
|
|
def residual_block_2d(
|
|
dim_in: int,
|
|
dim_out: int,
|
|
dim_hidden: int | None = None,
|
|
actvn: nn.Module | None = None,
|
|
norm_type: NormLayerName = "noop",
|
|
norm_num_groups: int = 8,
|
|
dilation: int = 1,
|
|
kernel_size: int = 3,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> ResidualBlock:
|
|
"""Create a simple 2D residual block."""
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if actvn is None:
|
|
actvn = nn.ReLU()
|
|
if dim_hidden is None:
|
|
dim_hidden = dim_out // 2
|
|
|
|
padding = (dilation * (kernel_size - 1)) // 2
|
|
|
|
def _create_block(d_in: int, d_out: int) -> list[nn.Module]:
|
|
layers: list[nn.Module] = [
|
|
norm_layer_2d(d_in, norm_type, num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations),
|
|
actvn,
|
|
operations.Conv2d(
|
|
d_in, d_out, kernel_size=kernel_size, stride=1,
|
|
dilation=dilation, padding=padding,
|
|
dtype=dtype, device=device,
|
|
),
|
|
]
|
|
return layers
|
|
|
|
residual = nn.Sequential(
|
|
*_create_block(dim_in, dim_hidden),
|
|
*_create_block(dim_hidden, dim_out),
|
|
)
|
|
shortcut = None
|
|
if dim_in != dim_out:
|
|
shortcut = operations.Conv2d(dim_in, dim_out, 1, dtype=dtype, device=device)
|
|
|
|
return ResidualBlock(residual, shortcut)
|
|
|
|
|
|
class FeatureFusionBlock2d(nn.Module):
|
|
"""Feature fusion for DPT."""
|
|
|
|
deconv: nn.Module
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in: int,
|
|
dim_out: int | None = None,
|
|
upsampling_mode: UpsamplingMode | None = None,
|
|
batch_norm: bool = False,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if dim_out is None:
|
|
dim_out = dim_in
|
|
self.resnet1 = self._residual_block(dim_in, batch_norm,
|
|
dtype=dtype, device=device, operations=operations)
|
|
self.resnet2 = self._residual_block(dim_in, batch_norm,
|
|
dtype=dtype, device=device, operations=operations)
|
|
|
|
if upsampling_mode is not None:
|
|
self.deconv = upsampling_layer(upsampling_mode, scale_factor=2, dim_in=dim_in,
|
|
dtype=dtype, device=device, operations=operations)
|
|
else:
|
|
self.deconv = nn.Sequential()
|
|
|
|
self.out_conv = operations.Conv2d(
|
|
dim_in, dim_out, kernel_size=1, stride=1, padding=0, bias=True,
|
|
dtype=dtype, device=device,
|
|
)
|
|
|
|
def forward(self, x0: torch.Tensor, x1: torch.Tensor | None = None) -> torch.Tensor:
|
|
x = x0
|
|
if x1 is not None:
|
|
res = self.resnet1(x1)
|
|
x = x + res
|
|
x = self.resnet2(x)
|
|
x = self.deconv(x)
|
|
x = self.out_conv(x)
|
|
return x
|
|
|
|
@staticmethod
|
|
def _residual_block(
|
|
num_features: int,
|
|
batch_norm: bool,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> ResidualBlock:
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
def _create_block(dim: int, bn: bool) -> list[nn.Module]:
|
|
layers: list[nn.Module] = [
|
|
nn.ReLU(False),
|
|
operations.Conv2d(
|
|
num_features, num_features, kernel_size=3, stride=1, padding=1,
|
|
bias=not bn, dtype=dtype, device=device,
|
|
),
|
|
]
|
|
if bn:
|
|
layers.append(nn.BatchNorm2d(dim))
|
|
return layers
|
|
|
|
residual = nn.Sequential(
|
|
*_create_block(dim=num_features, bn=batch_norm),
|
|
*_create_block(dim=num_features, bn=batch_norm),
|
|
)
|
|
return ResidualBlock(residual)
|
|
|
|
|
|
# ===========================================================================
|
|
# MultiresConvDecoder
|
|
# ===========================================================================
|
|
|
|
|
|
class MultiresConvDecoder(nn.Module):
|
|
"""Decoder for multi-resolution encodings."""
|
|
|
|
def __init__(
|
|
self,
|
|
dims_encoder: list[int],
|
|
dims_decoder: list[int] | int,
|
|
upsampling_mode: UpsamplingMode = "transposed_conv",
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.dims_encoder = list(dims_encoder)
|
|
|
|
if isinstance(dims_decoder, int):
|
|
self.dims_decoder = [dims_decoder] * len(self.dims_encoder)
|
|
else:
|
|
self.dims_decoder = list(dims_decoder)
|
|
|
|
if len(self.dims_decoder) != len(self.dims_encoder):
|
|
raise ValueError("Received dims_encoder and dims_decoder of different sizes.")
|
|
|
|
self.dim_out = self.dims_decoder[0]
|
|
num_encoders = len(self.dims_encoder)
|
|
|
|
conv0: nn.Module = (
|
|
operations.Conv2d(self.dims_encoder[0], self.dims_decoder[0],
|
|
kernel_size=1, bias=False, dtype=dtype, device=device)
|
|
if self.dims_encoder[0] != self.dims_decoder[0]
|
|
else nn.Identity()
|
|
)
|
|
|
|
convs: list[nn.Module] = [conv0]
|
|
for i in range(1, num_encoders):
|
|
convs.append(
|
|
operations.Conv2d(
|
|
self.dims_encoder[i], self.dims_decoder[i],
|
|
kernel_size=3, stride=1, padding=1, bias=False,
|
|
dtype=dtype, device=device,
|
|
)
|
|
)
|
|
self.convs = nn.ModuleList(convs)
|
|
|
|
fusions = []
|
|
for i in range(num_encoders):
|
|
fusions.append(
|
|
FeatureFusionBlock2d(
|
|
dim_in=self.dims_decoder[i],
|
|
dim_out=self.dims_decoder[i - 1] if i != 0 else self.dim_out,
|
|
upsampling_mode=upsampling_mode if i != 0 else None,
|
|
batch_norm=False,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
)
|
|
self.fusions = nn.ModuleList(fusions)
|
|
|
|
def forward(self, encodings: list[torch.Tensor]) -> torch.Tensor:
|
|
num_levels = len(encodings)
|
|
num_encoders = len(self.dims_encoder)
|
|
if num_levels != num_encoders:
|
|
raise ValueError(
|
|
f"Encoder output levels={num_levels} mismatch with expected={num_encoders}."
|
|
)
|
|
|
|
features = self.convs[-1](encodings[-1])
|
|
features = self.fusions[-1](features)
|
|
for i in range(num_levels - 2, -1, -1):
|
|
features_i = self.convs[i](encodings[i])
|
|
features = self.fusions[i](features, features_i)
|
|
return features
|
|
|
|
|
|
# ===========================================================================
|
|
# Normalizers
|
|
# ===========================================================================
|
|
|
|
|
|
class AffineRangeNormalizer(nn.Module):
|
|
"""Perform linear mapping from input_range to output_range."""
|
|
|
|
def __init__(
|
|
self,
|
|
input_range: tuple[float, float],
|
|
output_range: tuple[float, float] = (0, 1),
|
|
):
|
|
super().__init__()
|
|
input_min, input_max = input_range
|
|
output_min, output_max = output_range
|
|
if input_max <= input_min:
|
|
raise ValueError(f"Invalid input_range: {input_range}")
|
|
if output_max <= output_min:
|
|
raise ValueError(f"Invalid output_range: {output_range}")
|
|
self.scale = (output_max - output_min) / (input_max - input_min)
|
|
self.bias = output_min - input_min * self.scale
|
|
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
|
if self.scale != 1.0:
|
|
x = x * self.scale
|
|
if self.bias != 0.0:
|
|
x = x + self.bias
|
|
return x
|
|
|
|
|
|
class MeanStdNormalizer(nn.Module):
|
|
"""Normalizing image input by mean and std."""
|
|
|
|
mean: torch.Tensor
|
|
std_inv: torch.Tensor
|
|
|
|
def __init__(
|
|
self,
|
|
mean: Union[Sequence[float], torch.Tensor],
|
|
std: Union[Sequence[float], torch.Tensor],
|
|
):
|
|
super().__init__()
|
|
if not isinstance(mean, torch.Tensor):
|
|
mean = torch.as_tensor(mean).view(-1, 1, 1)
|
|
if not isinstance(std, torch.Tensor):
|
|
std = torch.as_tensor(std).view(-1, 1, 1)
|
|
self.register_buffer("mean", mean)
|
|
self.register_buffer("std_inv", 1.0 / std)
|
|
|
|
def forward(self, image: torch.Tensor) -> torch.Tensor:
|
|
return (image - self.mean) * self.std_inv
|
|
|
|
|
|
# ===========================================================================
|
|
# Monodepth
|
|
# ===========================================================================
|
|
|
|
|
|
class MonodepthDensePredictionTransformer(nn.Module):
|
|
"""Dense Prediction Transformer for monodepth."""
|
|
|
|
def __init__(
|
|
self,
|
|
encoder: SlidingPyramidNetwork,
|
|
decoder: MultiresConvDecoder,
|
|
last_dims: tuple[int, int],
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.normalizer = AffineRangeNormalizer(
|
|
input_range=(0, 1), output_range=(-1, 1)
|
|
)
|
|
self.encoder = encoder
|
|
self.decoder = decoder
|
|
|
|
dim_decoder = decoder.dim_out
|
|
self.head = nn.Sequential(
|
|
operations.Conv2d(dim_decoder, dim_decoder // 2, kernel_size=3, stride=1, padding=1,
|
|
dtype=dtype, device=device),
|
|
operations.ConvTranspose2d(
|
|
in_channels=dim_decoder // 2, out_channels=dim_decoder // 2,
|
|
kernel_size=2, stride=2, padding=0, bias=True,
|
|
dtype=dtype, device=device,
|
|
),
|
|
operations.Conv2d(
|
|
dim_decoder // 2, last_dims[0], kernel_size=3, stride=1, padding=1,
|
|
dtype=dtype, device=device,
|
|
),
|
|
nn.ReLU(True),
|
|
operations.Conv2d(last_dims[0], last_dims[1], kernel_size=1, stride=1, padding=0,
|
|
dtype=dtype, device=device),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
def forward(self, image: torch.Tensor) -> torch.Tensor:
|
|
encodings = self.encoder(self.normalizer(image))
|
|
num_encoder_features = len(self.encoder.dims_encoder)
|
|
features = self.decoder(encodings[:num_encoder_features])
|
|
disparity = self.head(features)
|
|
return disparity
|
|
|
|
def internal_resolution(self) -> int:
|
|
return self.encoder.internal_resolution()
|
|
|
|
|
|
class MonodepthWithEncodingAdaptor(nn.Module):
|
|
"""Monodepth model with feature maps."""
|
|
|
|
def __init__(
|
|
self,
|
|
monodepth_predictor: MonodepthDensePredictionTransformer,
|
|
return_encoder_features: bool,
|
|
return_decoder_features: bool,
|
|
num_monodepth_layers: int,
|
|
sorting_monodepth: bool,
|
|
):
|
|
super().__init__()
|
|
self.monodepth_predictor = monodepth_predictor
|
|
self.return_encoder_features = return_encoder_features
|
|
self.return_decoder_features = return_decoder_features
|
|
self.num_monodepth_layers = num_monodepth_layers
|
|
self.sorting_monodepth = sorting_monodepth
|
|
|
|
def forward(self, image: torch.Tensor) -> MonodepthOutput:
|
|
inputs = self.monodepth_predictor.normalizer(image)
|
|
encoder_output = self.monodepth_predictor.encoder(inputs)
|
|
|
|
num_encoder_features = len(self.monodepth_predictor.encoder.dims_encoder)
|
|
encoder_features = encoder_output[:num_encoder_features]
|
|
intermediate_features = encoder_output[num_encoder_features:]
|
|
decoder_features = self.monodepth_predictor.decoder(encoder_features)
|
|
disparity = self.monodepth_predictor.head(decoder_features)
|
|
|
|
if self.num_monodepth_layers == 2 and self.sorting_monodepth:
|
|
first_layer_disparity = disparity.max(dim=1, keepdims=True).values
|
|
second_layer_disparity = disparity.min(dim=1, keepdims=True).values
|
|
disparity = torch.cat([first_layer_disparity, second_layer_disparity], dim=1)
|
|
|
|
output_features: list[torch.Tensor] = []
|
|
if self.return_encoder_features:
|
|
output_features.extend(encoder_features)
|
|
if self.return_decoder_features:
|
|
output_features.append(decoder_features)
|
|
|
|
return MonodepthOutput(
|
|
disparity=disparity,
|
|
encoder_features=encoder_features,
|
|
decoder_features=decoder_features,
|
|
output_features=output_features,
|
|
intermediate_features=intermediate_features,
|
|
)
|
|
|
|
def get_feature_dims(self) -> list[int]:
|
|
dims: list[int] = []
|
|
if self.return_encoder_features:
|
|
dims.extend(self.monodepth_predictor.encoder.dims_encoder)
|
|
if self.return_decoder_features:
|
|
dims.append(self.monodepth_predictor.decoder.dim_out)
|
|
return dims
|
|
|
|
def internal_resolution(self) -> int:
|
|
return self.monodepth_predictor.internal_resolution()
|
|
|
|
def replicate_head(self, num_repeat: int):
|
|
"""Replicate the last convolution layer for multi-layer depth."""
|
|
conv_last = copy.deepcopy(self.monodepth_predictor.head[4])
|
|
self.monodepth_predictor.head[4].out_channels = num_repeat
|
|
self.monodepth_predictor.head[4].weight = nn.Parameter(
|
|
conv_last.weight.repeat(num_repeat, 1, 1, 1)
|
|
)
|
|
self.monodepth_predictor.head[4].bias = nn.Parameter(
|
|
conv_last.bias.repeat(num_repeat)
|
|
)
|
|
|
|
|
|
# ===========================================================================
|
|
# Gaussian Decoder
|
|
# ===========================================================================
|
|
|
|
|
|
def _create_project_upsample_block(
|
|
dim_in: int,
|
|
dim_out: int,
|
|
upsample_layers: int,
|
|
dim_intermediate: int | None = None,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> nn.Module:
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if dim_intermediate is None:
|
|
dim_intermediate = dim_out
|
|
blocks: list[nn.Module] = [
|
|
operations.Conv2d(
|
|
in_channels=dim_in, out_channels=dim_intermediate,
|
|
kernel_size=1, stride=1, padding=0, bias=False,
|
|
dtype=dtype, device=device,
|
|
)
|
|
]
|
|
blocks += [
|
|
operations.ConvTranspose2d(
|
|
in_channels=dim_intermediate if i == 0 else dim_out,
|
|
out_channels=dim_out, kernel_size=2, stride=2, padding=0, bias=False,
|
|
dtype=dtype, device=device,
|
|
)
|
|
for i in range(upsample_layers)
|
|
]
|
|
return nn.Sequential(*blocks)
|
|
|
|
|
|
class SkipConvBackbone(nn.Module):
|
|
"""A wrapper around a conv layer that behaves like a BaseBackbone."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in: int,
|
|
dim_out: int,
|
|
kernel_size: int,
|
|
stride_out: int,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.stride_out = stride_out
|
|
if stride_out == 1 and kernel_size != 1:
|
|
raise ValueError("We only support kernel_size = 1 if stride_out is 1.")
|
|
padding = (kernel_size - 1) // 2
|
|
self.conv = operations.Conv2d(
|
|
dim_in, dim_out, kernel_size=kernel_size, stride=stride_out, padding=padding,
|
|
dtype=dtype, device=device,
|
|
)
|
|
|
|
def forward(
|
|
self, input_features: torch.Tensor, encodings: list[torch.Tensor] | None = None
|
|
) -> ImageFeatures:
|
|
output = self.conv(input_features)
|
|
return ImageFeatures(texture_features=output, geometry_features=output)
|
|
|
|
@property
|
|
def stride(self) -> int:
|
|
return self.stride_out
|
|
|
|
|
|
class GaussianDensePredictionTransformer(nn.Module):
|
|
"""Dense Prediction Transformer for Gaussian features."""
|
|
|
|
norm_type: NormLayerName
|
|
|
|
def __init__(
|
|
self,
|
|
decoder: MultiresConvDecoder,
|
|
dim_in: int,
|
|
dim_out: int,
|
|
stride_out: int,
|
|
image_encoder_params, # GaussianDecoderParams
|
|
image_encoder_type: DPTImageEncoderType = "skip_conv",
|
|
norm_type: NormLayerName = "group_norm",
|
|
norm_num_groups: int = 8,
|
|
use_depth_input: bool = True,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
):
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.decoder = decoder
|
|
self.dim_in = dim_in
|
|
self.dim_out = dim_out
|
|
self.stride_out = stride_out
|
|
self.norm_type = norm_type
|
|
self.norm_num_groups = norm_num_groups
|
|
self.use_depth_input = use_depth_input
|
|
self.image_encoder_type = image_encoder_type
|
|
|
|
enc_dim_in = self.dim_in if use_depth_input else self.dim_in - 1
|
|
image_encoder_params.dim_in = enc_dim_in
|
|
image_encoder_params.dim_out = decoder.dim_out
|
|
self.image_encoder = self._create_image_encoder(
|
|
image_encoder_params, stride_out,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
|
|
self.fusion = FeatureFusionBlock2d(decoder.dim_out,
|
|
dtype=dtype, device=device, operations=operations)
|
|
|
|
if stride_out == 1:
|
|
self.upsample = _create_project_upsample_block(
|
|
decoder.dim_out, decoder.dim_out, upsample_layers=1,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
elif stride_out == 2:
|
|
self.upsample = nn.Identity()
|
|
else:
|
|
raise ValueError("We only support stride is 1 or 2 for DPT backbone.")
|
|
|
|
self.texture_head = self._create_head(
|
|
dim_decoder=decoder.dim_out, dim_out=self.dim_out,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.geometry_head = self._create_head(
|
|
dim_decoder=decoder.dim_out, dim_out=self.dim_out,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
|
|
def _create_head(
|
|
self, dim_decoder: int, dim_out: int,
|
|
dtype=None, device=None, operations=None,
|
|
) -> nn.Module:
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
return nn.Sequential(
|
|
residual_block_2d(
|
|
dim_in=dim_decoder, dim_out=dim_decoder, dim_hidden=dim_decoder // 2,
|
|
norm_type=self.norm_type, norm_num_groups=self.norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
),
|
|
residual_block_2d(
|
|
dim_in=dim_decoder, dim_hidden=dim_decoder // 2, dim_out=dim_decoder,
|
|
norm_type=self.norm_type, norm_num_groups=self.norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
),
|
|
nn.ReLU(),
|
|
operations.Conv2d(dim_decoder, dim_out, kernel_size=1, stride=1,
|
|
dtype=dtype, device=device),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
def _create_image_encoder(
|
|
self, image_encoder_params, stride_out: int,
|
|
dtype=None, device=None, operations=None,
|
|
) -> nn.Module:
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if self.image_encoder_type == "skip_conv":
|
|
return SkipConvBackbone(
|
|
image_encoder_params.dim_in, image_encoder_params.dim_out,
|
|
kernel_size=3 if stride_out != 1 else 1, stride_out=stride_out,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
elif self.image_encoder_type == "skip_conv_kernel2":
|
|
return SkipConvBackbone(
|
|
image_encoder_params.dim_in, image_encoder_params.dim_out,
|
|
kernel_size=stride_out, stride_out=stride_out,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unsupported image encoder type: {self.image_encoder_type}")
|
|
|
|
def forward(
|
|
self, input_features: torch.Tensor, encodings: list[torch.Tensor]
|
|
) -> ImageFeatures:
|
|
features = self.decoder(encodings).contiguous()
|
|
features = self.upsample(features)
|
|
|
|
if self.use_depth_input:
|
|
skip_features = self.image_encoder(input_features).texture_features
|
|
else:
|
|
skip_features = self.image_encoder(input_features[:, :3].contiguous())
|
|
features = self.fusion(features, skip_features)
|
|
|
|
texture_features = self.texture_head(features)
|
|
geometry_features = self.geometry_head(features)
|
|
|
|
return ImageFeatures(
|
|
texture_features=texture_features,
|
|
geometry_features=geometry_features,
|
|
)
|
|
|
|
@property
|
|
def stride(self) -> int:
|
|
return self.stride_out
|
|
|
|
|
|
# ===========================================================================
|
|
# Prediction Head
|
|
# ===========================================================================
|
|
|
|
|
|
class DirectPredictionHead(nn.Module):
|
|
"""Decodes features into delta values using convolutions."""
|
|
|
|
def __init__(
|
|
self,
|
|
feature_dim: int,
|
|
num_layers: int,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> None:
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self.num_layers = num_layers
|
|
self.geometry_prediction_head = operations.Conv2d(
|
|
feature_dim, 3 * num_layers, 1, dtype=dtype, device=device,
|
|
)
|
|
self.texture_prediction_head = operations.Conv2d(
|
|
feature_dim, (14 - 3) * num_layers, 1, dtype=dtype, device=device,
|
|
)
|
|
|
|
def forward(self, image_features: ImageFeatures) -> torch.Tensor:
|
|
delta_values_geometry = self.geometry_prediction_head(image_features.geometry_features)
|
|
delta_values_texture = self.texture_prediction_head(image_features.texture_features)
|
|
delta_values_geometry = delta_values_geometry.unflatten(1, (3, self.num_layers))
|
|
delta_values_texture = delta_values_texture.unflatten(1, (14 - 3, self.num_layers))
|
|
return torch.cat([delta_values_geometry, delta_values_texture], dim=1)
|
|
|
|
|
|
# ===========================================================================
|
|
# Initializer
|
|
# ===========================================================================
|
|
|
|
|
|
def _create_base_xy(
|
|
depth: torch.Tensor, stride: int, num_layers: int
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Create base x and y coordinates for Gaussians in NDC space."""
|
|
device = depth.device
|
|
batch_size, _, image_height, image_width = depth.shape
|
|
xx = torch.arange(0.5 * stride, image_width, stride, device=device)
|
|
yy = torch.arange(0.5 * stride, image_height, stride, device=device)
|
|
xx = 2 * xx / image_width - 1.0
|
|
yy = 2 * yy / image_height - 1.0
|
|
|
|
xx, yy = torch.meshgrid(xx, yy, indexing="xy")
|
|
base_x_ndc = xx[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
|
|
base_y_ndc = yy[None, None, None].repeat(batch_size, 1, num_layers, 1, 1)
|
|
return base_x_ndc, base_y_ndc
|
|
|
|
|
|
def _create_base_scale(disparity: torch.Tensor, disparity_scale_factor: float) -> torch.Tensor:
|
|
"""Create base scale for Gaussians."""
|
|
inverse_disparity = torch.ones_like(disparity) / disparity
|
|
return inverse_disparity * disparity_scale_factor
|
|
|
|
|
|
def _rescale_depth(
|
|
depth: torch.Tensor, depth_min: float = 1.0, depth_max: float = 1e2
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Rescale a depth image tensor."""
|
|
current_depth_min = depth.flatten(depth.ndim - 3).min(dim=-1).values
|
|
depth_factor = depth_min / (current_depth_min + 1e-6)
|
|
depth = (depth * depth_factor[..., None, None, None]).clamp(max=depth_max)
|
|
return depth, depth_factor
|
|
|
|
|
|
class MultiLayerInitializer(nn.Module):
|
|
"""Initialize Gaussians with multilayer representation."""
|
|
|
|
def __init__(
|
|
self,
|
|
num_layers: int,
|
|
stride: int,
|
|
base_depth: float,
|
|
scale_factor: float,
|
|
disparity_factor: float,
|
|
color_option: ColorInitOption = "first_layer",
|
|
first_layer_depth_option: DepthInitOption = "surface_min",
|
|
rest_layer_depth_option: DepthInitOption = "surface_min",
|
|
normalize_depth: bool = True,
|
|
feature_input_stop_grad: bool = True,
|
|
) -> None:
|
|
super().__init__()
|
|
self.num_layers = num_layers
|
|
self.stride = stride
|
|
self.base_depth = base_depth
|
|
self.scale_factor = scale_factor
|
|
self.disparity_factor = disparity_factor
|
|
self.color_option = color_option
|
|
self.first_layer_depth_option = first_layer_depth_option
|
|
self.rest_layer_depth_option = rest_layer_depth_option
|
|
self.normalize_depth = normalize_depth
|
|
self.feature_input_stop_grad = feature_input_stop_grad
|
|
|
|
def prepare_feature_input(self, image: torch.Tensor, depth: torch.Tensor) -> torch.Tensor:
|
|
if self.feature_input_stop_grad:
|
|
image = image.detach()
|
|
depth = depth.detach()
|
|
normalized_disparity = self.disparity_factor / depth
|
|
features_in = torch.cat([image, normalized_disparity], dim=1)
|
|
features_in = 2.0 * features_in - 1.0
|
|
return features_in
|
|
|
|
def forward(self, image: torch.Tensor, depth: torch.Tensor) -> InitializerOutput:
|
|
image = image.contiguous()
|
|
depth = depth.contiguous()
|
|
device = depth.device
|
|
batch_size, _, image_height, image_width = depth.shape
|
|
base_height, base_width = image_height // self.stride, image_width // self.stride
|
|
|
|
global_scale: torch.Tensor | None = None
|
|
if self.normalize_depth:
|
|
depth, depth_factor = _rescale_depth(depth)
|
|
global_scale = 1.0 / depth_factor
|
|
|
|
def _create_disparity_layers(n_layers: int = 1) -> torch.Tensor:
|
|
disparity = torch.linspace(1.0 / self.base_depth, 0.0, n_layers + 1, device=device)
|
|
return disparity[None, None, :-1, None, None].repeat(
|
|
batch_size, 1, 1, base_height, base_width
|
|
)
|
|
|
|
def _create_surface_layer(d: torch.Tensor, depth_pooling_mode: str) -> torch.Tensor:
|
|
disp = 1.0 / d
|
|
if depth_pooling_mode == "min":
|
|
disp = torch.max_pool2d(disp, self.stride, self.stride)
|
|
elif depth_pooling_mode == "max":
|
|
disp = -torch.max_pool2d(-disp, self.stride, self.stride)
|
|
else:
|
|
raise ValueError(f"Invalid depth pooling mode {depth_pooling_mode}.")
|
|
return disp[:, :, None, :, :]
|
|
|
|
if self.first_layer_depth_option == "surface_min":
|
|
first_disparity = _create_surface_layer(depth[:, 0:1], "min")
|
|
elif self.first_layer_depth_option == "surface_max":
|
|
first_disparity = _create_surface_layer(depth[:, 0:1], "max")
|
|
elif self.first_layer_depth_option in ("base_depth", "linear_disparity"):
|
|
first_disparity = _create_disparity_layers()
|
|
else:
|
|
raise ValueError(f"Unknown depth init option: {self.first_layer_depth_option}.")
|
|
|
|
if self.num_layers == 1:
|
|
disparity = first_disparity
|
|
else:
|
|
following_depth = depth if depth.shape[1] == 1 else depth[:, 1:]
|
|
if self.rest_layer_depth_option == "surface_min":
|
|
following_disparity = _create_surface_layer(following_depth, "min")
|
|
elif self.rest_layer_depth_option == "surface_max":
|
|
following_disparity = _create_surface_layer(following_depth, "max")
|
|
elif self.rest_layer_depth_option == "base_depth":
|
|
following_disparity = torch.cat(
|
|
[_create_disparity_layers() for _ in range(self.num_layers - 1)], dim=2,
|
|
)
|
|
elif self.rest_layer_depth_option == "linear_disparity":
|
|
following_disparity = _create_disparity_layers(self.num_layers - 1)
|
|
else:
|
|
raise ValueError(f"Unknown depth init option: {self.rest_layer_depth_option}.")
|
|
disparity = torch.cat([first_disparity, following_disparity], dim=2)
|
|
|
|
base_x_ndc, base_y_ndc = _create_base_xy(depth, self.stride, self.num_layers)
|
|
disparity_scale_factor = 2 * self.scale_factor * self.stride / float(image_width)
|
|
base_scales = _create_base_scale(disparity, disparity_scale_factor)
|
|
|
|
base_quaternions = torch.tensor([1.0, 0.0, 0.0, 0.0], device=device)
|
|
base_quaternions = base_quaternions[None, :, None, None, None]
|
|
base_opacities = torch.tensor([min(1.0 / self.num_layers, 0.5)], device=device)
|
|
base_colors = torch.empty(
|
|
batch_size, 3, self.num_layers, base_height, base_width, device=device
|
|
).fill_(0.5)
|
|
|
|
if self.color_option == "none":
|
|
pass
|
|
elif self.color_option == "first_layer":
|
|
base_colors[:, :, 0] = F.avg_pool2d(image, self.stride, self.stride)
|
|
elif self.color_option == "all_layers":
|
|
temp = F.avg_pool2d(image, self.stride, self.stride)
|
|
base_colors = temp[:, :, None, :, :].repeat(1, 1, self.num_layers, 1, 1)
|
|
else:
|
|
raise ValueError(f"Unknown color init option: {self.color_option}.")
|
|
|
|
features_in = self.prepare_feature_input(image, depth)
|
|
base_gaussians = GaussianBaseValues(
|
|
mean_x_ndc=base_x_ndc, mean_y_ndc=base_y_ndc,
|
|
mean_inverse_z_ndc=disparity, scales=base_scales,
|
|
quaternions=base_quaternions, colors=base_colors, opacities=base_opacities,
|
|
)
|
|
return InitializerOutput(
|
|
gaussian_base_values=base_gaussians,
|
|
feature_input=features_in,
|
|
global_scale=global_scale,
|
|
)
|
|
|
|
|
|
# ===========================================================================
|
|
# Gaussian Composer
|
|
# ===========================================================================
|
|
|
|
|
|
def _get_scale_activation_constant(max_scale: float, min_scale: float) -> tuple[float, float]:
|
|
constant_a = (max_scale - min_scale) / (1 - min_scale) / (max_scale - 1)
|
|
x = torch.tensor((1.0 - min_scale) / (max_scale - min_scale))
|
|
constant_b = torch.log(x / (1 - x)).item()
|
|
return constant_a, constant_b
|
|
|
|
|
|
class GaussianComposer(nn.Module):
|
|
"""Converts base values and deltas into Gaussians. Pure tensor ops, no learnable layers."""
|
|
|
|
color_activation_type: str
|
|
opacity_activation_type: str
|
|
|
|
def __init__(
|
|
self,
|
|
delta_factor, # DeltaFactor
|
|
min_scale: float,
|
|
max_scale: float,
|
|
color_activation_type: str,
|
|
opacity_activation_type: str,
|
|
color_space: ColorSpace,
|
|
base_scale_on_predicted_mean: bool,
|
|
scale_factor: int = 1,
|
|
) -> None:
|
|
super().__init__()
|
|
self.delta_factor = delta_factor
|
|
self.max_scale = max_scale
|
|
self.min_scale = min_scale
|
|
self.color_activation_type = color_activation_type
|
|
self.opacity_activation_type = opacity_activation_type
|
|
self.color_space = color_space
|
|
self.scale_factor = scale_factor
|
|
self.base_scale_on_predicted_mean = base_scale_on_predicted_mean
|
|
|
|
def upsample_delta_value(self, delta: torch.Tensor, scale_factor: int = 1):
|
|
B, C, N, H, W = delta.shape
|
|
new_height = H * scale_factor
|
|
new_width = W * scale_factor
|
|
return F.interpolate(
|
|
delta.view(B, C * N, H, W), scale_factor=scale_factor,
|
|
).view(B, C, N, new_height, new_width)
|
|
|
|
def forward(
|
|
self,
|
|
delta: torch.Tensor,
|
|
base_values: GaussianBaseValues,
|
|
global_scale: torch.Tensor | None = None,
|
|
flatten_output: bool = True,
|
|
) -> Gaussians3D:
|
|
scale_factor = self.scale_factor
|
|
actual_scale_factor = base_values.mean_x_ndc.shape[-1] // delta.shape[-1]
|
|
if scale_factor != 1 and actual_scale_factor != 1:
|
|
delta = self.upsample_delta_value(delta, scale_factor)
|
|
|
|
mean_vectors = self._forward_mean(base_values, delta)
|
|
base_scales = (
|
|
(base_values.scales * base_values.mean_inverse_z_ndc * mean_vectors[:, 2:3, ...])
|
|
if self.base_scale_on_predicted_mean
|
|
else base_values.scales
|
|
)
|
|
singular_values = self._scale_activation(
|
|
base_scales, delta[:, 3:6], self.min_scale, self.max_scale,
|
|
)
|
|
quaternions = self._quaternion_activation(base_values.quaternions, delta[:, 6:10])
|
|
colors = self._color_activation(base_values.colors, delta[:, 10:13])
|
|
opacities = self._opacity_activation(base_values.opacities, delta[:, 13])
|
|
|
|
if flatten_output:
|
|
mean_vectors = mean_vectors.permute(0, 2, 3, 4, 1).flatten(1, 3)
|
|
singular_values = singular_values.permute(0, 2, 3, 4, 1).flatten(1, 3)
|
|
quaternions = quaternions.permute(0, 2, 3, 4, 1).flatten(1, 3)
|
|
colors = colors.permute(0, 2, 3, 4, 1).flatten(1, 3)
|
|
opacities = opacities.flatten(1, 3)
|
|
|
|
if global_scale is not None:
|
|
mean_vectors = global_scale[:, None, None] * mean_vectors
|
|
singular_values = global_scale[:, None, None] * singular_values
|
|
|
|
return Gaussians3D(
|
|
mean_vectors=mean_vectors, singular_values=singular_values,
|
|
quaternions=quaternions, colors=colors, opacities=opacities,
|
|
)
|
|
|
|
def _forward_mean(self, base_values: GaussianBaseValues, delta: torch.Tensor) -> torch.Tensor:
|
|
delta_factor = torch.tensor(
|
|
[self.delta_factor.xy, self.delta_factor.xy, self.delta_factor.z],
|
|
device=delta.device,
|
|
)[None, :, None, None, None]
|
|
|
|
dtype = base_values.mean_x_ndc.dtype
|
|
device = base_values.mean_x_ndc.device
|
|
target_shape = (1, 3, 1, 1, 1)
|
|
mean_x_mask = torch.tensor([1.0, 0.0, 0.0], dtype=dtype, device=device).reshape(target_shape)
|
|
mean_y_mask = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device).reshape(target_shape)
|
|
mean_z_mask = torch.tensor([0.0, 0.0, 1.0], dtype=dtype, device=device).reshape(target_shape)
|
|
|
|
mean_vectors_ndc = (
|
|
base_values.mean_x_ndc.repeat(target_shape) * mean_x_mask
|
|
+ base_values.mean_y_ndc.repeat(target_shape) * mean_y_mask
|
|
+ base_values.mean_inverse_z_ndc.repeat(target_shape) * mean_z_mask
|
|
)
|
|
return self._mean_activation(mean_vectors_ndc, delta_factor * delta[:, :3])
|
|
|
|
def _mean_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
|
|
xx = base[:, 0:1] + learned_delta[:, 0:1]
|
|
yy = base[:, 1:2] + learned_delta[:, 1:2]
|
|
a = base[:, 2:3]
|
|
b = learned_delta[:, 2:3]
|
|
inverse_zz = F.softplus(_inverse_softplus(a) + b)
|
|
zz = 1.0 / (inverse_zz + 1e-3)
|
|
return torch.cat([zz * xx, zz * yy, zz], dim=1)
|
|
|
|
def _scale_activation(
|
|
self, base: torch.Tensor, learned_delta: torch.Tensor,
|
|
min_scale: float, max_scale: float,
|
|
) -> torch.Tensor:
|
|
constant_a, constant_b = _get_scale_activation_constant(max_scale, min_scale)
|
|
scale_factor = (max_scale - min_scale) * torch.sigmoid(
|
|
constant_a * self.delta_factor.scale * learned_delta + constant_b
|
|
) + min_scale
|
|
return base * scale_factor
|
|
|
|
def _quaternion_activation(
|
|
self, base: torch.Tensor, learned_delta: torch.Tensor
|
|
) -> torch.Tensor:
|
|
return base + self.delta_factor.quaternion * learned_delta
|
|
|
|
def _color_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
|
|
if self.color_activation_type == "sigmoid":
|
|
base = torch.clamp(base, min=0.01, max=0.99)
|
|
elif self.color_activation_type in ("exp", "softplus"):
|
|
base = torch.clamp(base, min=0.01)
|
|
fwd, inv = _ACTIVATIONS[self.color_activation_type]
|
|
colors = fwd(inv(base) + self.delta_factor.color * learned_delta)
|
|
if self.color_space == "linearRGB":
|
|
colors = sRGB2linearRGB(colors)
|
|
return colors
|
|
|
|
def _opacity_activation(self, base: torch.Tensor, learned_delta: torch.Tensor) -> torch.Tensor:
|
|
fwd, inv = _ACTIVATIONS[self.opacity_activation_type]
|
|
return fwd(inv(base) + self.delta_factor.opacity * learned_delta)
|
|
|
|
|
|
# ===========================================================================
|
|
# UNet Encoder / Decoder (used by LearnedAlignment)
|
|
# ===========================================================================
|
|
|
|
|
|
class UNetEncoder(nn.Module):
|
|
"""Encoder of UNet model."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_in: int,
|
|
width: list[int] | int,
|
|
steps: int = 6,
|
|
norm_type: NormLayerName = "group_norm",
|
|
norm_num_groups: int = 8,
|
|
blocks_per_layer: int = 2,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> None:
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if blocks_per_layer < 1:
|
|
raise ValueError("blocks_per_layer must be >= 1.")
|
|
|
|
self.dim_in = dim_in
|
|
self.width = width
|
|
self.num_steps = steps
|
|
|
|
self.convs_down = nn.ModuleList()
|
|
|
|
if isinstance(width, int):
|
|
self.output_dims = [width << i for i in range(0, steps + 1)]
|
|
else:
|
|
if len(width) != (steps + 1):
|
|
raise ValueError("Length of width should match the steps for UNetEncoder.")
|
|
self.output_dims = list(width)
|
|
|
|
self.conv_in = nn.Sequential(
|
|
operations.Conv2d(self.dim_in, self.output_dims[0], 3, stride=1, padding=1,
|
|
dtype=dtype, device=device),
|
|
norm_layer_2d(self.output_dims[0], norm_type, num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
for i_step in range(steps):
|
|
input_width = self.output_dims[i_step]
|
|
current_width = self.output_dims[i_step + 1]
|
|
convs_down_i = nn.Sequential(
|
|
nn.AvgPool2d(2, stride=2),
|
|
residual_block_2d(
|
|
input_width, current_width,
|
|
norm_type=norm_type, norm_num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
),
|
|
*[
|
|
residual_block_2d(
|
|
current_width, current_width,
|
|
norm_type=norm_type, norm_num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
for _ in range(blocks_per_layer - 1)
|
|
],
|
|
)
|
|
self.convs_down.append(convs_down_i)
|
|
|
|
def forward(self, input: torch.Tensor) -> list[torch.Tensor]:
|
|
features = []
|
|
feat_i = self.conv_in(input)
|
|
features.append(feat_i)
|
|
for conv_down in self.convs_down:
|
|
feat_i = conv_down(feat_i)
|
|
features.append(feat_i)
|
|
return features
|
|
|
|
@property
|
|
def out_width(self) -> int:
|
|
return self.output_dims[-1]
|
|
|
|
|
|
class UNetDecoder(nn.Module):
|
|
"""Decoder of UNet model."""
|
|
|
|
def __init__(
|
|
self,
|
|
dim_out: int,
|
|
width: list[int] | int,
|
|
steps: int = 5,
|
|
norm_type: NormLayerName = "group_norm",
|
|
norm_num_groups: int = 8,
|
|
blocks_per_layer: int = 2,
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> None:
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
if blocks_per_layer < 1:
|
|
raise ValueError("blocks_per_layer must be >= 1.")
|
|
|
|
self.dim_out = dim_out
|
|
self.convs_up = nn.ModuleList()
|
|
|
|
if isinstance(width, int):
|
|
self.input_dims = [width >> i for i in range(0, steps + 1)]
|
|
else:
|
|
self.input_dims = list(width)[::-1][: steps + 1]
|
|
|
|
for i_step in range(steps):
|
|
input_width = self.input_dims[i_step]
|
|
current_width = self.input_dims[i_step + 1]
|
|
convs_up_i = nn.Sequential(
|
|
nn.Upsample(scale_factor=2),
|
|
residual_block_2d(
|
|
input_width * (1 if i_step == 0 else 2),
|
|
current_width,
|
|
norm_type=norm_type, norm_num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
),
|
|
*[
|
|
residual_block_2d(
|
|
current_width, current_width,
|
|
norm_type=norm_type, norm_num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
for _ in range(blocks_per_layer - 1)
|
|
],
|
|
)
|
|
self.convs_up.append(convs_up_i)
|
|
|
|
last_width = self.input_dims[-1]
|
|
self.conv_out = nn.Sequential(
|
|
norm_layer_2d(last_width * 2, norm_type, num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations),
|
|
nn.ReLU(),
|
|
operations.Conv2d(last_width * 2, dim_out, 1, dtype=dtype, device=device),
|
|
norm_layer_2d(dim_out, norm_type, num_groups=norm_num_groups,
|
|
dtype=dtype, device=device, operations=operations),
|
|
nn.ReLU(),
|
|
)
|
|
|
|
def forward(self, features: list[torch.Tensor]) -> torch.Tensor:
|
|
i_feature_layer = len(features) - 1
|
|
out = self.convs_up[0](features[i_feature_layer])
|
|
i_feature_layer -= 1
|
|
for conv_up in self.convs_up[1:]:
|
|
out = conv_up(torch.cat([out, features[i_feature_layer]], dim=1))
|
|
i_feature_layer -= 1
|
|
out = self.conv_out(torch.cat([out, features[i_feature_layer]], dim=1))
|
|
return out
|
|
|
|
|
|
# ===========================================================================
|
|
# Learned Alignment
|
|
# ===========================================================================
|
|
|
|
|
|
class LearnedAlignment(nn.Module):
|
|
"""Aligns tensors using a UNet."""
|
|
|
|
def __init__(
|
|
self,
|
|
steps: int = 4,
|
|
stride: int = 8,
|
|
base_width: int = 16,
|
|
depth_decoder_features: bool = False,
|
|
depth_decoder_dim: int = 256,
|
|
activation_type: str = "exp",
|
|
dtype=None,
|
|
device=None,
|
|
operations=None,
|
|
) -> None:
|
|
super().__init__()
|
|
if operations is None:
|
|
operations = comfy.ops.manual_cast
|
|
self._act_fwd, self._act_inv = _ACTIVATIONS[activation_type]
|
|
bias_value = self._act_inv(torch.tensor(1.0))
|
|
|
|
self.depth_decoder_features = depth_decoder_features
|
|
dim_in = 2 + depth_decoder_dim if depth_decoder_features else 2
|
|
|
|
def is_power_of_two(n: int) -> bool:
|
|
if n <= 0:
|
|
return False
|
|
return (n & (n - 1)) == 0
|
|
|
|
if not is_power_of_two(stride):
|
|
raise ValueError(f"Stride {stride} is not a power of two.")
|
|
|
|
steps_decoder = steps - int(math.log2(stride))
|
|
if steps_decoder < 1:
|
|
raise ValueError(f"{steps_decoder} must be >= 1.")
|
|
widths = [min(base_width << i, 1024) for i in range(steps + 1)]
|
|
self.encoder = UNetEncoder(
|
|
dim_in=dim_in, width=widths, steps=steps, norm_num_groups=4,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.decoder = UNetDecoder(
|
|
dim_out=widths[0], width=widths, steps=steps_decoder, norm_num_groups=4,
|
|
dtype=dtype, device=device, operations=operations,
|
|
)
|
|
self.conv_out = operations.Conv2d(widths[0], 1, 1, bias=True,
|
|
dtype=dtype, device=device)
|
|
|
|
def forward(
|
|
self,
|
|
tensor_src: torch.Tensor,
|
|
tensor_tgt: torch.Tensor,
|
|
depth_decoder_features: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
tensor_src = 1.0 / tensor_src.clamp(min=1e-4)
|
|
tensor_tgt = 1.0 / tensor_tgt.clamp(min=1e-4)
|
|
tensor_input = torch.cat([tensor_src, tensor_tgt], dim=1)
|
|
if self.depth_decoder_features:
|
|
height, width = tensor_src.shape[-2:]
|
|
upsampled_encodings = F.interpolate(
|
|
depth_decoder_features, size=(height, width), mode="bilinear",
|
|
)
|
|
tensor_input = torch.cat([tensor_input, upsampled_encodings], dim=1)
|
|
features = self.encoder(tensor_input)
|
|
output = self.conv_out(self.decoder(features))
|
|
alignment_map_lowres = self._act_fwd(output)
|
|
if alignment_map_lowres.shape[-2:] != tensor_src.shape[-2:]:
|
|
alignment_map = F.interpolate(
|
|
alignment_map_lowres, size=tensor_src.shape[-2:],
|
|
mode="bilinear", align_corners=False,
|
|
)
|
|
else:
|
|
alignment_map = alignment_map_lowres
|
|
return alignment_map
|
|
|
|
|
|
# ===========================================================================
|
|
# Top-level model
|
|
# ===========================================================================
|
|
|
|
|
|
class DepthAlignment(nn.Module):
|
|
"""Depth alignment wrapper."""
|
|
|
|
def __init__(self, scale_map_estimator: nn.Module | None):
|
|
super().__init__()
|
|
self.scale_map_estimator = scale_map_estimator
|
|
|
|
def forward(
|
|
self,
|
|
monodepth: torch.Tensor,
|
|
depth: torch.Tensor,
|
|
depth_decoder_features: torch.Tensor | None = None,
|
|
):
|
|
if depth is not None and self.scale_map_estimator is not None:
|
|
depth_alignment_map = self.scale_map_estimator(
|
|
monodepth[:, 0:1], depth, depth_decoder_features
|
|
)
|
|
monodepth = depth_alignment_map * monodepth
|
|
else:
|
|
depth_alignment_map = torch.ones_like(monodepth)
|
|
return monodepth, depth_alignment_map
|
|
|
|
|
|
class RGBGaussianPredictor(nn.Module):
|
|
"""Predicts 3D Gaussians from images."""
|
|
|
|
feature_model: nn.Module
|
|
|
|
def __init__(
|
|
self,
|
|
init_model: nn.Module,
|
|
monodepth_model: MonodepthWithEncodingAdaptor,
|
|
feature_model: nn.Module,
|
|
prediction_head: nn.Module,
|
|
gaussian_composer: GaussianComposer,
|
|
scale_map_estimator: nn.Module | None,
|
|
dtype=None,
|
|
) -> None:
|
|
super().__init__()
|
|
self.dtype = dtype
|
|
self.init_model = init_model
|
|
self.feature_model = feature_model
|
|
self.monodepth_model = monodepth_model
|
|
self.prediction_head = prediction_head
|
|
self.gaussian_composer = gaussian_composer
|
|
self.depth_alignment = DepthAlignment(scale_map_estimator)
|
|
|
|
def encode(self, image: torch.Tensor):
|
|
monodepth_output = self.monodepth_model(image)
|
|
return monodepth_output, image
|
|
|
|
def decode(
|
|
self,
|
|
monodepth_output,
|
|
image: torch.Tensor,
|
|
disparity_factor: torch.Tensor,
|
|
depth: torch.Tensor | None = None,
|
|
) -> Gaussians3D:
|
|
monodepth_disparity = monodepth_output.disparity
|
|
disparity_factor = disparity_factor[:, None, None, None]
|
|
monodepth = disparity_factor / monodepth_disparity.clamp(min=1e-4, max=1e4)
|
|
|
|
monodepth, _ = self.depth_alignment(
|
|
monodepth, depth, monodepth_output.decoder_features,
|
|
)
|
|
|
|
init_output = self.init_model(image, monodepth)
|
|
image_features = self.feature_model(
|
|
init_output.feature_input, encodings=monodepth_output.output_features
|
|
)
|
|
delta_values = self.prediction_head(image_features)
|
|
return self.gaussian_composer(
|
|
delta=delta_values,
|
|
base_values=init_output.gaussian_base_values,
|
|
global_scale=init_output.global_scale,
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
image: torch.Tensor,
|
|
disparity_factor: torch.Tensor,
|
|
depth: torch.Tensor | None = None,
|
|
) -> Gaussians3D:
|
|
monodepth_output = self.monodepth_model(image)
|
|
monodepth_disparity = monodepth_output.disparity
|
|
|
|
disparity_factor = disparity_factor[:, None, None, None]
|
|
monodepth = disparity_factor / monodepth_disparity.clamp(min=1e-4, max=1e4)
|
|
|
|
monodepth, _ = self.depth_alignment(
|
|
monodepth, depth, monodepth_output.decoder_features,
|
|
)
|
|
|
|
init_output = self.init_model(image, monodepth)
|
|
image_features = self.feature_model(
|
|
init_output.feature_input, encodings=monodepth_output.output_features
|
|
)
|
|
delta_values = self.prediction_head(image_features)
|
|
return self.gaussian_composer(
|
|
delta=delta_values,
|
|
base_values=init_output.gaussian_base_values,
|
|
global_scale=init_output.global_scale,
|
|
)
|
|
|
|
def internal_resolution(self) -> int:
|
|
return self.monodepth_model.internal_resolution()
|
|
|
|
@property
|
|
def output_resolution(self) -> int:
|
|
return self.internal_resolution() // 2
|