(self, x, time_emb, guidance=None)
| 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): |