| 941 | |
| 942 | |
| 943 | class TextImageProjection(nn.Module): |
| 944 | def __init__( |
| 945 | self, |
| 946 | text_embed_dim: int = 1024, |
| 947 | image_embed_dim: int = 768, |
| 948 | cross_attention_dim: int = 768, |
| 949 | num_image_text_embeds: int = 10, |
| 950 | ): |
| 951 | super().__init__() |
| 952 | |
| 953 | self.num_image_text_embeds = num_image_text_embeds |
| 954 | self.image_embeds = nn.Linear(image_embed_dim, self.num_image_text_embeds * cross_attention_dim) |
| 955 | self.text_proj = nn.Linear(text_embed_dim, cross_attention_dim) |
| 956 | |
| 957 | def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): |
| 958 | batch_size = text_embeds.shape[0] |
| 959 | |
| 960 | # image |
| 961 | image_text_embeds = self.image_embeds(image_embeds) |
| 962 | image_text_embeds = image_text_embeds.reshape(batch_size, self.num_image_text_embeds, -1) |
| 963 | |
| 964 | # text |
| 965 | text_embeds = self.text_proj(text_embeds) |
| 966 | |
| 967 | return torch.cat([image_text_embeds, text_embeds], dim=1) |
| 968 | |
| 969 | |
| 970 | class ImageProjection(nn.Module): |
no outgoing calls
no test coverage detected