| 260 | return out |
| 261 | |
| 262 | class ImageProjModel(nn.Module): |
| 263 | def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): |
| 264 | super().__init__() |
| 265 | |
| 266 | self.cross_attention_dim = cross_attention_dim |
| 267 | self.clip_extra_context_tokens = clip_extra_context_tokens |
| 268 | self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim) |
| 269 | self.norm = nn.LayerNorm(cross_attention_dim) |
| 270 | |
| 271 | def forward(self, image_embeds): |
| 272 | embeds = image_embeds |
| 273 | x = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim) |
| 274 | x = self.norm(x) |
| 275 | return x |