| 910 | |
| 911 | |
| 912 | def patchify(x, patch_size): |
| 913 | if patch_size == 1: |
| 914 | return x |
| 915 | |
| 916 | if x.dim() != 5: |
| 917 | raise ValueError(f"Invalid input shape: {x.shape}") |
| 918 | # x shape: [batch_size, channels, frames, height, width] |
| 919 | batch_size, channels, frames, height, width = x.shape |
| 920 | |
| 921 | # Ensure height and width are divisible by patch_size |
| 922 | if height % patch_size != 0 or width % patch_size != 0: |
| 923 | raise ValueError(f"Height ({height}) and width ({width}) must be divisible by patch_size ({patch_size})") |
| 924 | |
| 925 | # Reshape to [batch_size, channels, frames, height//patch_size, patch_size, width//patch_size, patch_size] |
| 926 | x = x.view(batch_size, channels, frames, height // patch_size, patch_size, width // patch_size, patch_size) |
| 927 | |
| 928 | # Rearrange to [batch_size, channels * patch_size * patch_size, frames, height//patch_size, width//patch_size] |
| 929 | x = x.permute(0, 1, 6, 4, 2, 3, 5).contiguous() |
| 930 | x = x.view(batch_size, channels * patch_size * patch_size, frames, height // patch_size, width // patch_size) |
| 931 | |
| 932 | return x |
| 933 | |
| 934 | |
| 935 | def unpatchify(x, patch_size): |