| 477 | |
| 478 | |
| 479 | class Blur(Module): |
| 480 | def __init__(self): |
| 481 | super().__init__() |
| 482 | f = torch.Tensor([1, 2, 1]) |
| 483 | self.register_buffer("f", f) |
| 484 | |
| 485 | def forward(self, x, space_only=False, time_only=False): |
| 486 | assert not (space_only and time_only) |
| 487 | |
| 488 | f = self.f |
| 489 | |
| 490 | if space_only: |
| 491 | f = einsum("i, j -> i j", f, f) |
| 492 | f = rearrange(f, "... -> 1 1 ...") |
| 493 | elif time_only: |
| 494 | f = rearrange(f, "f -> 1 f 1 1") |
| 495 | else: |
| 496 | f = einsum("i, j, k -> i j k", f, f, f) |
| 497 | f = rearrange(f, "... -> 1 ...") |
| 498 | |
| 499 | is_images = x.ndim == 4 |
| 500 | |
| 501 | if is_images: |
| 502 | x = rearrange(x, "b c h w -> b c 1 h w") |
| 503 | |
| 504 | out = filter3d(x, f, normalized=True) |
| 505 | |
| 506 | if is_images: |
| 507 | out = rearrange(out, "b c 1 h w -> b c h w") |
| 508 | |
| 509 | return out |
| 510 | |
| 511 | |
| 512 | class DiscriminatorBlock(Module): |