| 211 | |
| 212 | |
| 213 | class TimestepEmbedding_(nn.Module): |
| 214 | def __init__(self, n_time=1024, n_emb=320, n_out=1280) -> None: |
| 215 | super().__init__() |
| 216 | self.emb = nn.Embedding(n_time, n_emb) |
| 217 | self.f_1 = nn.Linear(n_emb, n_out) |
| 218 | self.f_2 = nn.Linear(n_out, n_out) |
| 219 | |
| 220 | def forward(self, x) -> torch.Tensor: |
| 221 | x = self.emb(x) |
| 222 | x = self.f_1(x) |
| 223 | x = F.silu(x) |
| 224 | return self.f_2(x) |
| 225 | |
| 226 | |
| 227 | class ImageEmbedding(nn.Module): |