(self, x, timesteps, low_res=None, **kwargs)
| 745 | super().__init__(image_size, in_channels * 2, *args, **kwargs) |
| 746 | |
| 747 | def forward(self, x, timesteps, low_res=None, **kwargs): |
| 748 | _, _, new_height, new_width = x.shape |
| 749 | upsampled = F.interpolate(low_res, (new_height, new_width), mode="bilinear") |
| 750 | x = th.cat([x, upsampled], dim=1) |
| 751 | return super().forward(x, timesteps, **kwargs) |
| 752 | |
| 753 | |
| 754 | class EncoderUNetModel(nn.Module): |