(self, channels, padding=1, extra_padding=False)
| 69 | |
| 70 | class DownSampler(torch.nn.Module): |
| 71 | def __init__(self, channels, padding=1, extra_padding=False): |
| 72 | super().__init__() |
| 73 | self.conv = torch.nn.Conv2d(channels, channels, 3, stride=2, padding=padding) |
| 74 | self.extra_padding = extra_padding |
| 75 | |
| 76 | def forward(self, hidden_states, time_emb, text_emb, res_stack, **kwargs): |
| 77 | if self.extra_padding: |