(
self,
in_channels: int,
time_embed_dim: int,
act_fn: str = "silu",
out_dim: int = None,
post_act_fn: Optional[str] = None,
cond_proj_dim=None,
sample_proj_bias=True,
)
| 716 | |
| 717 | class TimestepEmbedding(nn.Module): |
| 718 | def __init__( |
| 719 | self, |
| 720 | in_channels: int, |
| 721 | time_embed_dim: int, |
| 722 | act_fn: str = "silu", |
| 723 | out_dim: int = None, |
| 724 | post_act_fn: Optional[str] = None, |
| 725 | cond_proj_dim=None, |
| 726 | sample_proj_bias=True, |
| 727 | ): |
| 728 | super().__init__() |
| 729 | |
| 730 | self.linear_1 = nn.Linear(in_channels, time_embed_dim, sample_proj_bias) |
| 731 | |
| 732 | if cond_proj_dim is not None: |
| 733 | self.cond_proj = nn.Linear(cond_proj_dim, in_channels, bias=False) |
| 734 | else: |
| 735 | self.cond_proj = None |
| 736 | |
| 737 | self.act = get_activation(act_fn) |
| 738 | |
| 739 | if out_dim is not None: |
| 740 | time_embed_dim_out = out_dim |
| 741 | else: |
| 742 | time_embed_dim_out = time_embed_dim |
| 743 | self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim_out, sample_proj_bias) |
| 744 | |
| 745 | if post_act_fn is None: |
| 746 | self.post_act = None |
| 747 | else: |
| 748 | self.post_act = get_activation(post_act_fn) |
| 749 | |
| 750 | def forward(self, sample, condition=None): |
| 751 | if condition is not None: |
nothing calls this directly
no test coverage detected