| 219 | |
| 220 | |
| 221 | class UpBlock2D(nn.Module): |
| 222 | |
| 223 | def __init__( |
| 224 | self, |
| 225 | in_channels, |
| 226 | prev_output_channel, |
| 227 | out_channels, |
| 228 | temb_channels, |
| 229 | num_layers=1, |
| 230 | resnet_eps=1e-6, |
| 231 | resnet_time_scale_shift="default", |
| 232 | resnet_act_fn="swish", |
| 233 | resnet_groups=32, |
| 234 | add_upsample=True, |
| 235 | ): |
| 236 | super().__init__() |
| 237 | resnets = [] |
| 238 | |
| 239 | for i in range(num_layers): |
| 240 | res_skip_channels = in_channels if (i == num_layers - |
| 241 | 1) else out_channels |
| 242 | resnet_in_channels = prev_output_channel if i == 0 else out_channels |
| 243 | |
| 244 | resnets.append( |
| 245 | ResnetBlock2D( |
| 246 | in_channels=resnet_in_channels + res_skip_channels, |
| 247 | out_channels=out_channels, |
| 248 | temb_channels=temb_channels, |
| 249 | eps=resnet_eps, |
| 250 | groups=resnet_groups, |
| 251 | time_embedding_norm=resnet_time_scale_shift, |
| 252 | )) |
| 253 | |
| 254 | self.resnets = nn.ModuleList(resnets) |
| 255 | self.upsamplers = None |
| 256 | if add_upsample: |
| 257 | self.upsamplers = nn.ModuleList([Upsample2D(out_channels)]) |
| 258 | |
| 259 | def forward(self, hidden_states, res_hidden_states_tuple, temb=None): |
| 260 | for resnet in self.resnets: |
| 261 | res_hidden_states = res_hidden_states_tuple[-1] |
| 262 | res_hidden_states_tuple = res_hidden_states_tuple[:-1] |
| 263 | hidden_states = torch.cat([hidden_states, res_hidden_states], |
| 264 | dim=1) |
| 265 | |
| 266 | hidden_states = resnet(hidden_states, temb) |
| 267 | |
| 268 | if self.upsamplers is not None: |
| 269 | for upsampler in self.upsamplers: |
| 270 | hidden_states = upsampler(hidden_states) |
| 271 | |
| 272 | return hidden_states |
| 273 | |
| 274 | |
| 275 | class CrossAttnDownBlock2D(nn.Module): |