| 1486 | |
| 1487 | class TextImageProjection(nn.Module): |
| 1488 | def __init__( |
| 1489 | self, |
| 1490 | text_embed_dim: int = 1024, |
| 1491 | image_embed_dim: int = 768, |
| 1492 | cross_attention_dim: int = 768, |
| 1493 | num_image_text_embeds: int = 10, |
| 1494 | ): |
| 1495 | super().__init__() |
| 1496 | |
| 1497 | self.num_image_text_embeds = num_image_text_embeds |
| 1498 | self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim) |
| 1499 | self.text_proj = nn.Linear(text_embed_dim, cross_attention_dim) |
| 1500 | |
| 1501 | def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): |
| 1502 | batch_size = text_embeds.shape[0] |