| 74 | return features |
| 75 | |
| 76 | class SplatDecoder(nn.Module): |
| 77 | def __init__(self, opt: Options, **kwargs): |
| 78 | super().__init__() |
| 79 | |
| 80 | self.opt = opt |
| 81 | self.width = opt.decoder_hidden_dim |
| 82 | self.patch_size = opt.patch_size |
| 83 | self.input_res = opt.down_resolution |
| 84 | self.num_layers = opt.decoder_num_layers |
| 85 | |
| 86 | if len(opt.down_resolution) > 0: |
| 87 | self.actual_input_res = opt.down_resolution |
| 88 | else: |
| 89 | self.actual_input_res = (opt.image_height, opt.image_width) |
| 90 | |
| 91 | self.transformer_decoder = TransformerConditionalDecoder( |
| 92 | input_res=self.actual_input_res, |
| 93 | patch_size=self.patch_size, |
| 94 | layers=self.num_layers, |
| 95 | width=self.width, |
| 96 | heads=self.width // 64, |
| 97 | window_size=opt.bwindow_size, |
| 98 | condition_dim=opt.hidden_dim, |
| 99 | condition_len=576 if opt.use_dino else 2304, |
| 100 | encoder_dim=opt.hidden_dim, |
| 101 | drop_path_rate=opt.drop_path_rate, |
| 102 | ) |
| 103 | self.token_len = (self.actual_input_res[0] // self.patch_size) * (self.actual_input_res[1] // self.patch_size) |
| 104 | |
| 105 | self.transformer_decoder.set_grad_checkpointing(opt.checkpointing) |
| 106 | |
| 107 | def forward(self, latent, condition=None): |
| 108 | features = self.transformer_decoder(latent, condition) # [B, V, N, D] |
| 109 | |
| 110 | return features |
| 111 | |
| 112 | inverse_sigmoid = lambda x: np.log(x / (1 - x)) |
| 113 | artanh = lambda x: 0.5 * np.log((1 + x) / (1 - x)) |