| 933 | |
| 934 | |
| 935 | def unpatchify(x, patch_size): |
| 936 | if patch_size == 1: |
| 937 | return x |
| 938 | |
| 939 | if x.dim() != 5: |
| 940 | raise ValueError(f"Invalid input shape: {x.shape}") |
| 941 | # x shape: [batch_size, (channels * patch_size * patch_size), frame, height, width] |
| 942 | batch_size, c_patches, frames, height, width = x.shape |
| 943 | channels = c_patches // (patch_size * patch_size) |
| 944 | |
| 945 | # Reshape to [b, c, patch_size, patch_size, f, h, w] |
| 946 | x = x.view(batch_size, channels, patch_size, patch_size, frames, height, width) |
| 947 | |
| 948 | # Rearrange to [b, c, f, h * patch_size, w * patch_size] |
| 949 | x = x.permute(0, 1, 4, 5, 3, 6, 2).contiguous() |
| 950 | x = x.view(batch_size, channels, frames, height * patch_size, width * patch_size) |
| 951 | |
| 952 | return x |
| 953 | |
| 954 | |
| 955 | class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin): |