A UNetModel that performs super-resolution. Expects an extra kwarg `low_res` to condition on a low-resolution image.
| 735 | |
| 736 | |
| 737 | class SuperResModel(UNetModel): |
| 738 | """ |
| 739 | A UNetModel that performs super-resolution. |
| 740 | |
| 741 | Expects an extra kwarg `low_res` to condition on a low-resolution image. |
| 742 | """ |
| 743 | |
| 744 | def __init__(self, image_size, in_channels, *args, **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): |
nothing calls this directly
no outgoing calls
no test coverage detected