| 1513 | |
| 1514 | class ImageProjection(nn.Module): |
| 1515 | def __init__( |
| 1516 | self, |
| 1517 | image_embed_dim: int = 768, |
| 1518 | cross_attention_dim: int = 768, |
| 1519 | num_image_text_embeds: int = 32, |
| 1520 | ): |
| 1521 | super().__init__() |
| 1522 | |
| 1523 | self.num_image_text_embeds = num_image_text_embeds |
| 1524 | self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim) |
| 1525 | self.norm = nn.LayerNorm(cross_attention_dim) |
| 1526 | |
| 1527 | def forward(self, image_embeds: torch.Tensor): |
| 1528 | batch_size = image_embeds.shape[0] |