r"""A 2D K-downsampling layer. Parameters: pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use.
| 254 | |
| 255 | # downsample/upsample layer used in k-upscaler, might be able to use FirDownsample2D/DirUpsample2D instead |
| 256 | class KDownsample2D(nn.Module): |
| 257 | r"""A 2D K-downsampling layer. |
| 258 | |
| 259 | Parameters: |
| 260 | pad_mode (`str`, *optional*, default to `"reflect"`): the padding mode to use. |
| 261 | """ |
| 262 | |
| 263 | def __init__(self, pad_mode: str = "reflect"): |
| 264 | super().__init__() |
| 265 | self.pad_mode = pad_mode |
| 266 | kernel_1d = torch.tensor([[1 / 8, 3 / 8, 3 / 8, 1 / 8]]) |
| 267 | self.pad = kernel_1d.shape[1] // 2 - 1 |
| 268 | self.register_buffer("kernel", kernel_1d.T @ kernel_1d, persistent=False) |
| 269 | |
| 270 | def forward(self, inputs: torch.Tensor) -> torch.Tensor: |
| 271 | inputs = F.pad(inputs, (self.pad,) * 4, self.pad_mode) |
| 272 | weight = inputs.new_zeros( |
| 273 | [ |
| 274 | inputs.shape[1], |
| 275 | inputs.shape[1], |
| 276 | self.kernel.shape[0], |
| 277 | self.kernel.shape[1], |
| 278 | ] |
| 279 | ) |
| 280 | indices = torch.arange(inputs.shape[1], device=inputs.device) |
| 281 | kernel = self.kernel.to(weight)[None, :].expand(inputs.shape[1], -1, -1) |
| 282 | weight[indices, indices] = kernel |
| 283 | return F.conv2d(inputs, weight, stride=2) |
| 284 | |
| 285 | |
| 286 | class CogVideoXDownsample3D(nn.Module): |