| 150 | |
| 151 | |
| 152 | class MidBlock(nn.Module): |
| 153 | def __init__( |
| 154 | self, |
| 155 | in_channels: int, |
| 156 | num_layers: int = 1, |
| 157 | attention: bool = True, |
| 158 | attention_heads: int = 16, |
| 159 | skip_scale: float = 1, |
| 160 | ): |
| 161 | super().__init__() |
| 162 | |
| 163 | nets = [] |
| 164 | attns = [] |
| 165 | # first layer |
| 166 | nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale)) |
| 167 | # more layers |
| 168 | for i in range(num_layers): |
| 169 | nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale)) |
| 170 | if attention: |
| 171 | attns.append(MVAttention(in_channels, attention_heads, skip_scale=skip_scale)) |
| 172 | else: |
| 173 | attns.append(None) |
| 174 | self.nets = nn.ModuleList(nets) |
| 175 | self.attns = nn.ModuleList(attns) |
| 176 | |
| 177 | def forward(self, x): |
| 178 | x = self.nets[0](x) |
| 179 | for attn, net in zip(self.attns, self.nets[1:]): |
| 180 | if attn: |
| 181 | x = attn(x) |
| 182 | x = net(x) |
| 183 | return x |
| 184 | |
| 185 | |
| 186 | class UpBlock(nn.Module): |