| 1232 | |
| 1233 | |
| 1234 | class TextImageTimeEmbedding(nn.Module): |
| 1235 | def __init__(self, text_embed_dim: int = 768, image_embed_dim: int = 768, time_embed_dim: int = 1536): |
| 1236 | super().__init__() |
| 1237 | self.text_proj = nn.Linear(text_embed_dim, time_embed_dim) |
| 1238 | self.text_norm = nn.LayerNorm(time_embed_dim) |
| 1239 | self.image_proj = nn.Linear(image_embed_dim, time_embed_dim) |
| 1240 | |
| 1241 | def forward(self, text_embeds: torch.Tensor, image_embeds: torch.Tensor): |
| 1242 | # text |
| 1243 | time_text_embeds = self.text_proj(text_embeds) |
| 1244 | time_text_embeds = self.text_norm(time_text_embeds) |
| 1245 | |
| 1246 | # image |
| 1247 | time_image_embeds = self.image_proj(image_embeds) |
| 1248 | |
| 1249 | return time_image_embeds + time_text_embeds |
| 1250 | |
| 1251 | |
| 1252 | class ImageTimeEmbedding(nn.Module): |
no outgoing calls
no test coverage detected