| 35 | self.pos_embed = torch.nn.Parameter(torch.zeros(1, self.pos_embed_max_size, self.pos_embed_max_size, embed_dim)) |
| 36 | |
| 37 | def cropped_pos_embed(self, height, width): |
| 38 | height = height // self.patch_size |
| 39 | width = width // self.patch_size |
| 40 | top = (self.pos_embed_max_size - height) // 2 |
| 41 | left = (self.pos_embed_max_size - width) // 2 |
| 42 | spatial_pos_embed = self.pos_embed[:, top : top + height, left : left + width, :].flatten(1, 2) |
| 43 | return spatial_pos_embed |
| 44 | |
| 45 | def forward(self, latent): |
| 46 | height, width = latent.shape[-2:] |