| 1 | import torch |
| 2 | from einops import rearrange |
| 3 | from torch import nn |
| 4 | |
| 5 | |
| 6 | def patchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: |
| 7 | """ |
| 8 | Rearrange spatial dimensions into channels. Divides image into patch_size x patch_size blocks |
| 9 | and moves pixels from each block into separate channels (space-to-depth). |
| 10 | Args: |
| 11 | x: Input tensor (4D or 5D) |
| 12 | patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, divides HxW into 4x4 blocks. |
| 13 | patch_size_t: Temporal patch size for frames. Default=1 (no temporal patching). |
| 14 | For 5D: (B, C, F, H, W) -> (B, Cx(patch_size_hw^2)x(patch_size_t), F/patch_size_t, H/patch_size_hw, W/patch_size_hw) |
| 15 | Example: (B, 3, 33, 512, 512) with patch_size_hw=4, patch_size_t=1 -> (B, 48, 33, 128, 128) |
| 16 | """ |
| 17 | if patch_size_hw == 1 and patch_size_t == 1: |
| 18 | return x |
| 19 | if x.dim() == 4: |
| 20 | x = rearrange(x, "b c (h q) (w r) -> b (c r q) h w", q=patch_size_hw, r=patch_size_hw) |
| 21 | elif x.dim() == 5: |
| 22 | x = rearrange( |
| 23 | x, |
| 24 | "b c (f p) (h q) (w r) -> b (c p r q) f h w", |
| 25 | p=patch_size_t, |
| 26 | q=patch_size_hw, |
| 27 | r=patch_size_hw, |
| 28 | ) |
| 29 | else: |
| 30 | raise ValueError(f"Invalid input shape: {x.shape}") |
| 31 | |
| 32 | return x |
| 33 | |
| 34 | |
| 35 | def unpatchify(x: torch.Tensor, patch_size_hw: int, patch_size_t: int = 1) -> torch.Tensor: |
| 36 | """ |
| 37 | Rearrange channels back into spatial dimensions. Inverse of patchify - moves pixels from |
| 38 | channels back into patch_size x patch_size blocks (depth-to-space). |
| 39 | Args: |
| 40 | x: Input tensor (4D or 5D) |
| 41 | patch_size_hw: Spatial patch size for height and width. With patch_size_hw=4, expands HxW by 4x. |
| 42 | patch_size_t: Temporal patch size for frames. Default=1 (no temporal expansion). |
| 43 | For 5D: (B, Cx(patch_size_hw^2)x(patch_size_t), F, H, W) -> (B, C, Fxpatch_size_t, Hxpatch_size_hw, Wxpatch_size_hw) |
| 44 | Example: (B, 48, 33, 128, 128) with patch_size_hw=4, patch_size_t=1 -> (B, 3, 33, 512, 512) |
| 45 | """ |
| 46 | if patch_size_hw == 1 and patch_size_t == 1: |
| 47 | return x |
| 48 | |
| 49 | if x.dim() == 4: |
| 50 | x = rearrange(x, "b (c r q) h w -> b c (h q) (w r)", q=patch_size_hw, r=patch_size_hw) |
| 51 | elif x.dim() == 5: |
| 52 | x = rearrange( |
| 53 | x, |
| 54 | "b (c p r q) f h w -> b c (f p) (h q) (w r)", |
| 55 | p=patch_size_t, |
| 56 | q=patch_size_hw, |
| 57 | r=patch_size_hw, |
| 58 | ) |
| 59 | |
| 60 | return x |
| 61 | |
| 62 | |
| 63 | class PerChannelStatistics(nn.Module): |
| 64 | """ |
| 65 | Per-channel statistics for normalizing and denormalizing the latent representation. |
| 66 | This statics is computed over the entire dataset and stored in model's checkpoint under VAE state_dict. |
| 67 | """ |
| 68 | |
| 69 | def __init__(self, latent_channels: int = 128): |
| 70 | super().__init__() |
| 71 | self.register_buffer("std-of-means", torch.empty(latent_channels)) |
| 72 | self.register_buffer("mean-of-means", torch.empty(latent_channels)) |
| 73 | |
| 74 | def un_normalize(self, x: torch.Tensor) -> torch.Tensor: |
| 75 | return (x * self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x)) + self.get_buffer("mean-of-means").view( |
| 76 | 1, -1, 1, 1, 1 |
| 77 | ).to(x) |
| 78 | |
| 79 | def normalize(self, x: torch.Tensor) -> torch.Tensor: |
| 80 | return (x - self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x)) / self.get_buffer("std-of-means").view( |
| 81 | 1, -1, 1, 1, 1 |
| 82 | ).to(x) |
| 83 |