(
self,
dim,
dim_out,
time_emb_dim=None,
dropout=0,
norm_groups=32,
attn_guide=False,
)
| 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) |
nothing calls this directly
no test coverage detected