| 177 | |
| 178 | |
| 179 | class ResnetBlocWithAttn(nn.Module): |
| 180 | def __init__( |
| 181 | self, |
| 182 | dim, |
| 183 | dim_out, |
| 184 | *, |
| 185 | time_emb_dim=None, |
| 186 | norm_groups=32, |
| 187 | dropout=0, |
| 188 | with_attn=False, |
| 189 | attn_guide=False, |
| 190 | ): |
| 191 | super().__init__() |
| 192 | self.with_attn = with_attn |
| 193 | self.res_block = ResnetBlock( |
| 194 | dim, |
| 195 | dim_out, |
| 196 | time_emb_dim, |
| 197 | norm_groups=norm_groups, |
| 198 | dropout=dropout, |
| 199 | attn_guide=attn_guide, |
| 200 | ) |
| 201 | if with_attn: |
| 202 | self.attn = SelfAttention(dim_out, norm_groups=norm_groups, nhead=8) |
| 203 | |
| 204 | def forward(self, x, time_emb, guidance=None): |
| 205 | x = self.res_block(x, time_emb, guidance) |
| 206 | if self.with_attn: |
| 207 | x = self.attn(x) |
| 208 | return x |
| 209 | |
| 210 | |
| 211 | class UNet(nn.Module): |