unisharp/models/feature_gaussian_decoder.py

from __future__ import annotations

import logging
from dataclasses import dataclass
from typing import NamedTuple

import torch
from torch import nn
from torch.nn import functional as F

from unisharp.models.blocks import FeatureFusionBlock2d, NormLayerName, residual_block_2d
from unisharp.models.decoder import MultiresConvDecoder


LOGGER = logging.getLogger(__name__)


class ImageFeatures(NamedTuple):

    texture_features: torch.Tensor
    geometry_features: torch.Tensor


class CircularAwareConv2d(nn.Conv2d):

    circular_horizontal: bool

    @classmethod
    def from_conv2d(cls, conv: nn.Conv2d) -> "CircularAwareConv2d":
        out = cls(
            in_channels=conv.in_channels,
            out_channels=conv.out_channels,
            kernel_size=conv.kernel_size,
            stride=conv.stride,
            padding=conv.padding,
            dilation=conv.dilation,
            groups=conv.groups,
            bias=conv.bias is not None,
            padding_mode=conv.padding_mode,
            device=conv.weight.device,
            dtype=conv.weight.dtype,
        )
        with torch.no_grad():
            out.weight.copy_(conv.weight)
            if conv.bias is not None and out.bias is not None:
                out.bias.copy_(conv.bias)
        return out

    def __init__(self, *args, **kwargs) -> None:  # type: ignore[no-untyped-def]
        super().__init__(*args, **kwargs)
        self.circular_horizontal = False

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        if not bool(self.circular_horizontal):
            return super().forward(input)
        pad_h, pad_w = self._padding_hw()
        if pad_h == 0 and pad_w == 0:
            return F.conv2d(
                input,
                self.weight,
                self.bias,
                self.stride,
                0,
                self.dilation,
                self.groups,
            )
        x = input
        if pad_w > 0:
            x = F.pad(x, (pad_w, pad_w, 0, 0), mode="circular")
        if pad_h > 0:
            x = F.pad(x, (0, 0, pad_h, pad_h), mode="constant", value=0.0)
        return F.conv2d(
            x,
            self.weight,
            self.bias,
            self.stride,
            0,
            self.dilation,
            self.groups,
        )

    def _padding_hw(self) -> tuple[int, int]:
        if isinstance(self.padding, tuple):
            if len(self.padding) == 2:
                return int(self.padding[0]), int(self.padding[1])
            return int(self.padding[0]), int(self.padding[0])
        return int(self.padding), int(self.padding)


def _convert_conv2d_modules_to_circular_aware(module: nn.Module) -> None:
    for name, child in list(module.named_children()):
        if isinstance(child, nn.Conv2d) and not isinstance(child, CircularAwareConv2d):
            setattr(module, name, CircularAwareConv2d.from_conv2d(child))
        else:
            _convert_conv2d_modules_to_circular_aware(child)


def _set_circular_horizontal(module: nn.Module, enabled: bool) -> None:
    for child in module.modules():
        if hasattr(child, "circular_horizontal"):
            child.circular_horizontal = bool(enabled)


@dataclass
class FeatureGaussianDecoderParams:

    dims_3d_in: tuple[int, int, int, int] = (128, 256, 512, 512)
    dims_3d_out: tuple[int, int, int, int] = (256, 512, 1024, 1024)
    
    dim_2d_in: int = 1024
    dim_2d_out: int = 256
    
    dim_decoder_out: int = 256
    
    dim_texture_out: int = 32
    dim_geometry_out: int = 32
    norm_type: NormLayerName = "group_norm"
    norm_num_groups: int = 8
    
    stride_out: int = 2

    use_learned_upsampling: bool = False
    target_resolution: tuple[int, int] | None = None


class Feature2DEncoder(nn.Module):
    
    def __init__(
        self,
        dim_in: int = 1024,
        dim_out: int = 256,
    ):
        super().__init__()
        
        self.process = nn.Sequential(
            nn.Conv2d(dim_in, 512, kernel_size=3, padding=1),
            nn.GroupNorm(8, 512),
            nn.GELU(),
            
            nn.Conv2d(512, dim_out, kernel_size=3, padding=1),
            nn.GroupNorm(8, dim_out),
            nn.GELU(),
        )
        
        self.dim_out = dim_out
    
    def forward(
        self,
        x: torch.Tensor,
        target_h: int,
        target_w: int,
    ) -> torch.Tensor:
        x = self.process(x)
        x = torch.nn.functional.interpolate(
            x,
            size=(target_h, target_w),
            mode="bilinear",
            align_corners=False,
        )
        
        return x


class Feature3DProjector(nn.Module):
    
    def __init__(
        self,
        dims_in: list[int],
        dims_out: list[int],
    ):
        super().__init__()
        
        if len(dims_in) != len(dims_out):
            raise ValueError(
                f"dims_in and dims_out must have same length, "
                f"got {len(dims_in)} vs {len(dims_out)}"
            )
        
        self.projectors = nn.ModuleList([
            nn.Conv2d(dim_in, dim_out, kernel_size=1, bias=False)
            for dim_in, dim_out in zip(dims_in, dims_out)
        ])
        
        self.dims_out = dims_out
        self.num_levels = len(dims_in)
    
    def forward(self, pyramid_features: list[torch.Tensor]) -> list[torch.Tensor]:
        if len(pyramid_features) != self.num_levels:
            raise ValueError(
                f"Expected {self.num_levels} pyramid features, got {len(pyramid_features)}"
            )
        
        return [proj(feat) for proj, feat in zip(self.projectors, pyramid_features)]


class LearnedUpsampler(nn.Module):
    
    def __init__(self, in_channels: int, out_channels: int):
        super().__init__()
        
        self.conv1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1)
        self.bn1 = nn.GroupNorm(8, in_channels)
        self.relu = nn.ReLU(inplace=True)
        
        self.conv2 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.bn2 = nn.GroupNorm(8, out_channels)
        
        self.skip = nn.Identity() if in_channels == out_channels else nn.Conv2d(in_channels, out_channels, kernel_size=1)
    
    def forward(self, x: torch.Tensor, target_h: int, target_w: int) -> torch.Tensor:
        x_upsampled = torch.nn.functional.interpolate(
            x, size=(target_h, target_w), mode="bilinear", align_corners=False
        )
        
        identity = self.skip(x_upsampled)
        
        out = self.conv1(x_upsampled)
        out = self.bn1(out)
        out = self.relu(out)
        
        out = self.conv2(out)
        out = self.bn2(out)
        
        out = out + identity
        out = self.relu(out)
        
        return out


def _create_project_upsample_block(dim_in: int, dim_out: int, upsample_layers: int) -> nn.Module:
    blocks: list[nn.Module] = [
        nn.Conv2d(
            in_channels=dim_in,
            out_channels=dim_out,
            kernel_size=1,
            stride=1,
            padding=0,
            bias=False,
        )
    ]
    blocks.extend(
        nn.Sequential(
            nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
            CircularAwareConv2d(
                in_channels=dim_out,
                out_channels=dim_out,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            ),
        )
        for _ in range(int(upsample_layers))
    )
    return nn.Sequential(*blocks)


class FeatureGaussianDecoder(nn.Module):
    
    def __init__(
        self,
        params: FeatureGaussianDecoderParams,
    ):
        super().__init__()
        
        self.params = params
        self.stride_out = params.stride_out
        self.norm_type = params.norm_type
        self.norm_num_groups = int(params.norm_num_groups)
        if int(self.stride_out) not in (1, 2):
            raise ValueError(f"FeatureGaussianDecoder only supports stride_out 1 or 2, got {self.stride_out}")
        
        self.feature_3d_projector = Feature3DProjector(
            dims_in=list(params.dims_3d_in),
            dims_out=list(params.dims_3d_out),
        )
        
        self.decoder = MultiresConvDecoder(
            dims_encoder=list(params.dims_3d_out),
            dims_decoder=params.dim_decoder_out,
        )
        if int(self.stride_out) == 1:
            self.upsample = _create_project_upsample_block(
                params.dim_decoder_out,
                params.dim_decoder_out,
                upsample_layers=1,
            )
        else:
            self.upsample = nn.Identity()
        
        self.feature_2d_encoder = Feature2DEncoder(
            dim_in=params.dim_2d_in,
            dim_out=params.dim_2d_out,
        )
        
        self.fusion = FeatureFusionBlock2d(
            params.dim_decoder_out,
            params.dim_2d_out,
        )
        
        self.texture_head = self._create_head(
            params.dim_decoder_out,
            params.dim_texture_out,
        )
        
        self.geometry_head = self._create_head(
            params.dim_decoder_out,
            params.dim_geometry_out,
        )

        if int(params.dim_2d_out) != int(params.dim_decoder_out):
            raise ValueError(
                "FeatureFusionBlock2d requires 2D skip channels to match decoder channels, "
                f"got dim_2d_out={params.dim_2d_out}, dim_decoder_out={params.dim_decoder_out}"
            )

        self.dim_out = params.dim_texture_out
        
        self.fused_upsampler = None

        _convert_conv2d_modules_to_circular_aware(self)
        
    def _create_head(self, dim_in: int, dim_out: int) -> nn.Module:
        return nn.Sequential(
            residual_block_2d(
                dim_in=dim_in,
                dim_out=dim_in,
                dim_hidden=dim_in // 2,
                norm_type=self.norm_type,
                norm_num_groups=self.norm_num_groups,
            ),
            residual_block_2d(
                dim_in=dim_in,
                dim_hidden=dim_in // 2,
                dim_out=dim_in,
                norm_type=self.norm_type,
                norm_num_groups=self.norm_num_groups,
            ),
            nn.ReLU(),
            nn.Conv2d(dim_in, dim_out, kernel_size=1, stride=1),
            nn.ReLU(),
        )
    
    def forward(
        self,
        features_2d: torch.Tensor,
        features_3d_pyramid: list[torch.Tensor],
        *,
        circular_horizontal: bool = False,
        target_hw: tuple[int, int] | None = None,
    ) -> ImageFeatures:
        _set_circular_horizontal(self, bool(circular_horizontal))
        features_3d_sorted = sorted(
            features_3d_pyramid,
            key=lambda t: int(t.shape[-2]) * int(t.shape[-1]),
            reverse=True
        )
        
        pyramid_projected = self.feature_3d_projector(features_3d_sorted)
        
        decoder_out = self.decoder(pyramid_projected).contiguous()
        decoder_out = self.upsample(decoder_out).contiguous()
        if target_hw is not None:
            target_h, target_w = int(target_hw[0]), int(target_hw[1])
            if target_h <= 0 or target_w <= 0:
                raise ValueError(f"target_hw must be positive, got {target_hw}")
            if tuple(decoder_out.shape[-2:]) != (target_h, target_w):
                decoder_out = F.interpolate(
                    decoder_out,
                    size=(target_h, target_w),
                    mode="bilinear",
                    align_corners=False,
                ).contiguous()
        
        target_h, target_w = decoder_out.shape[-2:]
        features_2d_proj = self.feature_2d_encoder(
            features_2d,
            target_h=target_h,
            target_w=target_w,
        )
        
        fused = self.fusion(decoder_out, features_2d_proj)
        
        if target_hw is not None:
            target_h, target_w = int(target_hw[0]), int(target_hw[1])
            if tuple(fused.shape[-2:]) != (target_h, target_w):
                raise RuntimeError(
                    "Feature decoder grid must match base Gaussian grid before heads, "
                    f"got fused={tuple(fused.shape[-2:])} target={(target_h, target_w)}. "
                    "Only high-channel decoder features may be adapted before fusion."
                )

        texture_features = self.texture_head(fused)
        geometry_features = self.geometry_head(fused)
        
        return ImageFeatures(
            texture_features=texture_features,
            geometry_features=geometry_features,
        )
    
    @property
    def stride(self) -> int:
        return self.stride_out


def create_feature_gaussian_decoder(
    params: FeatureGaussianDecoderParams | None = None,
) -> FeatureGaussianDecoder:
    if params is None:
        params = FeatureGaussianDecoderParams()
    
    return FeatureGaussianDecoder(params)