| 1216 | |
| 1217 | |
| 1218 | class TextTimeEmbedding(nn.Module): |
| 1219 | def __init__(self, encoder_dim: int, time_embed_dim: int, num_heads: int = 64): |
| 1220 | super().__init__() |
| 1221 | self.norm1 = nn.LayerNorm(encoder_dim) |
| 1222 | self.pool = AttentionPooling(num_heads, encoder_dim) |
| 1223 | self.proj = nn.Linear(encoder_dim, time_embed_dim) |
| 1224 | self.norm2 = nn.LayerNorm(time_embed_dim) |
| 1225 | |
| 1226 | def forward(self, hidden_states): |
| 1227 | hidden_states = self.norm1(hidden_states) |
| 1228 | hidden_states = self.pool(hidden_states) |
| 1229 | hidden_states = self.proj(hidden_states) |
| 1230 | hidden_states = self.norm2(hidden_states) |
| 1231 | return hidden_states |
| 1232 | |
| 1233 | |
| 1234 | class TextImageTimeEmbedding(nn.Module): |
no outgoing calls
no test coverage detected