(
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
)
| 162 | |
| 163 | class TimestepEmbedding(nn.Module): |
| 164 | def __init__( |
| 165 | self, |
| 166 | in_channels: int, |
| 167 | time_embed_dim: int, |
| 168 | act_fn: str = "silu", |
| 169 | out_dim: int = None, |
| 170 | post_act_fn: Optional[str] = None, |
| 171 | cond_proj_dim=None, |
| 172 | sample_proj_bias=True |
| 173 | ): |
| 174 | super().__init__() |
| 175 | linear_cls = nn.Linear |
| 176 | |
| 177 | self.linear_1 = linear_cls( |
| 178 | in_channels, |
| 179 | time_embed_dim, |
| 180 | bias=sample_proj_bias, |
| 181 | ) |
| 182 | |
| 183 | if cond_proj_dim is not None: |
| 184 | self.cond_proj = linear_cls( |
| 185 | cond_proj_dim, |
| 186 | in_channels, |
| 187 | bias=False, |
| 188 | ) |
| 189 | else: |
| 190 | self.cond_proj = None |
| 191 | |
| 192 | self.act = get_activation(act_fn) |
| 193 | |
| 194 | if out_dim is not None: |
| 195 | time_embed_dim_out = out_dim |
| 196 | else: |
| 197 | time_embed_dim_out = time_embed_dim |
| 198 | |
| 199 | self.linear_2 = linear_cls( |
| 200 | time_embed_dim, |
| 201 | time_embed_dim_out, |
| 202 | bias=sample_proj_bias, |
| 203 | ) |
| 204 | |
| 205 | if post_act_fn is None: |
| 206 | self.post_act = None |
| 207 | else: |
| 208 | self.post_act = get_activation(post_act_fn) |
| 209 | |
| 210 | def forward(self, sample, condition=None): |
| 211 | if condition is not None: |
nothing calls this directly
no test coverage detected