| 110 | |
| 111 | |
| 112 | class ResnetBlock(nn.Module): |
| 113 | def __init__( |
| 114 | self, |
| 115 | dim, |
| 116 | dim_out, |
| 117 | time_emb_dim=None, |
| 118 | dropout=0, |
| 119 | norm_groups=32, |
| 120 | attn_guide=False, |
| 121 | ): |
| 122 | super().__init__() |
| 123 | self.mlp = ( |
| 124 | nn.Sequential(Swish(), nn.Linear(time_emb_dim, dim_out)) |
| 125 | if exists(time_emb_dim) |
| 126 | else None |
| 127 | ) |
| 128 | |
| 129 | self.block1 = Block(dim, dim_out, groups=norm_groups) |
| 130 | self.block2 = Block(dim_out, dim_out, groups=norm_groups, dropout=dropout) |
| 131 | self.res_conv = nn.Conv2d(dim, dim_out, 1) if dim != dim_out else nn.Identity() |
| 132 | self.atten_guide = AttentiveGuide(dim_out) if attn_guide else nn.Identity() |
| 133 | |
| 134 | def forward(self, x, time_emb, guidance=None): |
| 135 | h = self.block1(x) |
| 136 | if exists(self.mlp): |
| 137 | h += self.mlp(time_emb)[:, :, None, None] |
| 138 | if exists(guidance): |
| 139 | # guidance should have the same shape as h |
| 140 | h = self.atten_guide(h, guidance) |
| 141 | h = self.block2(h) |
| 142 | return h + self.res_conv(x) |
| 143 | |
| 144 | |
| 145 | class SelfAttention(nn.Module): |