| 1563 | |
| 1564 | |
| 1565 | class CombinedTimestepLabelEmbeddings(nn.Module): |
| 1566 | def __init__(self, num_classes, embedding_dim, class_dropout_prob=0.1): |
| 1567 | super().__init__() |
| 1568 | |
| 1569 | self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=1) |
| 1570 | self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) |
| 1571 | self.class_embedder = LabelEmbedding(num_classes, embedding_dim, class_dropout_prob) |
| 1572 | |
| 1573 | def forward(self, timestep, class_labels, hidden_dtype=None): |
| 1574 | timesteps_proj = self.time_proj(timestep) |
| 1575 | timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype)) # (N, D) |
| 1576 | |
| 1577 | class_labels = self.class_embedder(class_labels) # (N, D) |
| 1578 | |
| 1579 | conditioning = timesteps_emb + class_labels # (N, D) |
| 1580 | |
| 1581 | return conditioning |
| 1582 | |
| 1583 | |
| 1584 | class CombinedTimestepTextProjEmbeddings(nn.Module): |