| 42 | init.constant_(m.bias, 0) |
| 43 | |
| 44 | class CondEmbedding(nn.Module): |
| 45 | def __init__(self, in_channels: int, cond_embed_dim: int, act_fn: str = "silu"): |
| 46 | super().__init__() |
| 47 | |
| 48 | self.linear = nn.Linear(in_channels, cond_embed_dim) |
| 49 | self.act = None |
| 50 | if act_fn == "silu": |
| 51 | self.act = nn.SiLU() |
| 52 | elif act_fn == "mish": |
| 53 | self.act = nn.Mish() |
| 54 | |
| 55 | def forward(self, sample): |
| 56 | sample = self.linear(sample) |
| 57 | |
| 58 | if self.act is not None: |
| 59 | sample = self.act(sample) |
| 60 | |
| 61 | return sample |
| 62 | |
| 63 | class UNet3DConditionModel(ModelMixin, ConfigMixin): |
| 64 | _supports_gradient_checkpointing = True |