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 | |
| class PositionNet(nn.Module): | |
| def __init__(self, in_dim, out_dim, fourier_freqs=8): | |
| super().__init__() | |
| self.in_dim = in_dim | |
| self.out_dim = out_dim | |
| self.fourier_embedder = FourierEmbedder(num_freqs=fourier_freqs) | |
| self.position_dim = fourier_freqs*2*4 # 2 is sin&cos, 4 is xyxy | |
| # -------------------------------------------------------------- # | |
| self.linears_text = nn.Sequential( | |
| nn.Linear( self.in_dim + self.position_dim, 512), | |
| nn.SiLU(), | |
| nn.Linear( 512, 512), | |
| nn.SiLU(), | |
| nn.Linear(512, out_dim), | |
| ) | |
| self.linears_image = nn.Sequential( | |
| nn.Linear( self.in_dim + self.position_dim, 512), | |
| nn.SiLU(), | |
| nn.Linear( 512, 512), | |
| nn.SiLU(), | |
| nn.Linear(512, out_dim), | |
| ) | |
| # -------------------------------------------------------------- # | |
| self.null_text_feature = torch.nn.Parameter(torch.zeros([self.in_dim])) | |
| self.null_image_feature = torch.nn.Parameter(torch.zeros([self.in_dim])) | |
| self.null_position_feature = torch.nn.Parameter(torch.zeros([self.position_dim])) | |
| def forward(self, boxes, masks, text_masks, image_masks, text_embeddings, image_embeddings): | |
| B, N, _ = boxes.shape | |
| masks = masks.unsqueeze(-1) # B*N*1 | |
| text_masks = text_masks.unsqueeze(-1) # B*N*1 | |
| image_masks = image_masks.unsqueeze(-1) # B*N*1 | |
| # embedding position (it may includes padding as placeholder) | |
| xyxy_embedding = self.fourier_embedder(boxes) # B*N*4 --> B*N*C | |
| # learnable null embedding | |
| text_null = self.null_text_feature.view(1,1,-1) # 1*1*C | |
| image_null = self.null_image_feature.view(1,1,-1) # 1*1*C | |
| xyxy_null = self.null_position_feature.view(1,1,-1) # 1*1*C | |
| # replace padding with learnable null embedding | |
| text_embeddings = text_embeddings*text_masks + (1-text_masks)*text_null | |
| image_embeddings = image_embeddings*image_masks + (1-image_masks)*image_null | |
| xyxy_embedding = xyxy_embedding*masks + (1-masks)*xyxy_null | |
| objs_text = self.linears_text( torch.cat([text_embeddings, xyxy_embedding], dim=-1) ) | |
| objs_image = self.linears_image( torch.cat([image_embeddings,xyxy_embedding], dim=-1) ) | |
| objs = torch.cat( [objs_text,objs_image], dim=1 ) | |
| assert objs.shape == torch.Size([B,N*2,self.out_dim]) | |
| return objs | |