| 510 | |
| 511 | |
| 512 | class TextImageTimeEmbedding(nn.Module): |
| 513 | def __init__(self, text_embed_dim: int = 768, image_embed_dim: int = 768, time_embed_dim: int = 1536): |
| 514 | super().__init__() |
| 515 | self.text_proj = nn.Linear(text_embed_dim, time_embed_dim) |
| 516 | self.text_norm = nn.LayerNorm(time_embed_dim) |
| 517 | self.image_proj = nn.Linear(image_embed_dim, time_embed_dim) |
| 518 | |
| 519 | def forward(self, text_embeds: torch.FloatTensor, image_embeds: torch.FloatTensor): |
| 520 | # text |
| 521 | time_text_embeds = self.text_proj(text_embeds) |
| 522 | time_text_embeds = self.text_norm(time_text_embeds) |
| 523 | |
| 524 | # image |
| 525 | time_image_embeds = self.image_proj(image_embeds) |
| 526 | |
| 527 | return time_image_embeds + time_text_embeds |
| 528 | |
| 529 | |
| 530 | class ImageTimeEmbedding(nn.Module): |
no outgoing calls
no test coverage detected