| 479 | class TimestepEmbedding(Module): |
| 480 | |
| 481 | def __init__(self, |
| 482 | in_channels: int, |
| 483 | time_embed_dim: int, |
| 484 | act_fn: str = "silu", |
| 485 | out_dim: int = None, |
| 486 | post_act_fn: Optional[str] = None, |
| 487 | cond_proj_dim=None, |
| 488 | sample_proj_bias=True, |
| 489 | mapping=None, |
| 490 | dtype=None): |
| 491 | super().__init__() |
| 492 | tp_group = mapping.tp_group |
| 493 | tp_size = mapping.tp_size |
| 494 | self.linear_1 = ColumnLinear(in_channels, |
| 495 | time_embed_dim, |
| 496 | sample_proj_bias, |
| 497 | tp_group=tp_group, |
| 498 | tp_size=tp_size, |
| 499 | dtype=dtype, |
| 500 | gather_output=False) |
| 501 | |
| 502 | if cond_proj_dim is not None: |
| 503 | self.cond_proj = Linear(cond_proj_dim, |
| 504 | in_channels, |
| 505 | bias=False, |
| 506 | dtype=dtype) |
| 507 | else: |
| 508 | self.cond_proj = None |
| 509 | |
| 510 | self.act = ACT2FN[act_fn] |
| 511 | |
| 512 | if out_dim is not None: |
| 513 | time_embed_dim_out = out_dim |
| 514 | else: |
| 515 | time_embed_dim_out = time_embed_dim |
| 516 | self.linear_2 = RowLinear(time_embed_dim, |
| 517 | time_embed_dim_out, |
| 518 | sample_proj_bias, |
| 519 | tp_group=tp_group, |
| 520 | tp_size=tp_size, |
| 521 | dtype=dtype) |
| 522 | |
| 523 | if post_act_fn is None: |
| 524 | self.post_act = None |
| 525 | else: |
| 526 | self.post_act = ACT2FN[post_act_fn] |
| 527 | |
| 528 | def forward(self, sample, condition=None): |
| 529 | if condition is not None: |