Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| from ldm.modules.attention import BasicTransformerBlock | |
| from ldm.modules.diffusionmodules.util import checkpoint, FourierEmbedder | |
| import torch.nn.functional as F | |
| from ..attention import SelfAttention, FeedForward | |
| from .convnext import convnext_tiny | |
| class PositionNet(nn.Module): | |
| def __init__(self, resize_input=448, out_dim=768): | |
| super().__init__() | |
| self.resize_input = resize_input | |
| self.down_factor = 32 # determined by the convnext backbone | |
| self.out_dim = out_dim | |
| assert self.resize_input % self.down_factor == 0 | |
| self.convnext_tiny_backbone = convnext_tiny(pretrained=True) | |
| self.num_tokens = (self.resize_input // self.down_factor) ** 2 | |
| convnext_feature_dim = 768 | |
| self.pos_embedding = nn.Parameter(torch.empty(1, self.num_tokens, convnext_feature_dim).normal_(std=0.02)) # from BERT | |
| self.linears = nn.Sequential( | |
| nn.Linear( convnext_feature_dim, 512), | |
| nn.SiLU(), | |
| nn.Linear( 512, 512), | |
| nn.SiLU(), | |
| nn.Linear(512, out_dim), | |
| ) | |
| self.null_feature = torch.nn.Parameter(torch.zeros([convnext_feature_dim])) | |
| def forward(self, normal, mask): | |
| B = normal.shape[0] | |
| # token from edge map | |
| normal = torch.nn.functional.interpolate(normal, self.resize_input) | |
| normal_feature = self.convnext_tiny_backbone(normal) | |
| objs = normal_feature.reshape(B, -1, self.num_tokens) | |
| objs = objs.permute(0, 2, 1) # N*Num_tokens*dim | |
| # expand null token | |
| null_objs = self.null_feature.view(1,1,-1) | |
| null_objs = null_objs.repeat(B,self.num_tokens,1) | |
| # mask replacing | |
| mask = mask.view(-1,1,1) | |
| objs = objs*mask + null_objs*(1-mask) | |
| # add pos | |
| objs = objs + self.pos_embedding | |
| # fuse them | |
| objs = self.linears(objs) | |
| assert objs.shape == torch.Size([B,self.num_tokens,self.out_dim]) | |
| return objs | |