| 81 | |
| 82 | |
| 83 | class UpSampler(torch.nn.Module): |
| 84 | def __init__(self, channels): |
| 85 | super().__init__() |
| 86 | self.conv = torch.nn.Conv2d(channels, channels, 3, padding=1) |
| 87 | |
| 88 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 89 | hidden_states = torch.nn.functional.interpolate(hidden_states, scale_factor=2.0, mode="nearest") |
| 90 | hidden_states = self.conv(hidden_states) |
| 91 | return hidden_states, time_emb, text_emb, res_stack |
| 92 | |
| 93 | |
| 94 | class ResnetBlock(torch.nn.Module): |