facok
0d82606fb0
Video VAEs output [B, C, T, H, W]. Patchify now merges T into batch, processes all frames, and unpatchify restores the 5D shape. Updated all callers to pass extra_shape through.
46 lines
1.7 KiB
Python
46 lines
1.7 KiB
Python
"""Patchify/unpatchify for FLUX-family latent tensors (patch_size=2, auto-detect channels)."""
|
||
|
||
from einops import rearrange
|
||
|
||
|
||
def patchify(x):
|
||
"""Convert latent [B, C, H, W] or [B, C, T, H, W] → patch sequence [B, L, C*4].
|
||
|
||
For video VAEs (5D input with T frames), merges T into the batch dimension
|
||
so all frames are patchified together. The original shape is returned as
|
||
extra_shape for unpatchify to restore.
|
||
|
||
L = (H/2) * (W/2), d = C * 2 * 2.
|
||
Returns (patches, h_len, w_len, extra_shape) where extra_shape is None for 4D
|
||
or (B_orig, C, T) for 5D.
|
||
"""
|
||
extra_shape = None
|
||
if x.ndim == 5:
|
||
B_orig, C, T, H, W = x.shape
|
||
extra_shape = (B_orig, C, T)
|
||
# Merge B and T: [B*T, C, H, W]
|
||
x = x.permute(0, 2, 1, 3, 4).reshape(B_orig * T, C, H, W)
|
||
B, C, H, W = x.shape
|
||
h_len = H // 2
|
||
w_len = W // 2
|
||
patches = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2)
|
||
return patches, h_len, w_len, extra_shape
|
||
|
||
|
||
def unpatchify(patches, h_len, w_len, extra_shape=None):
|
||
"""Convert patch sequence [B, L, C*4] → latent [B, C, H, W] or [B, C, T, H, W].
|
||
|
||
Auto-detects channel count from patch dimension: C = D / 4.
|
||
If extra_shape is provided, restores the 5D video format.
|
||
"""
|
||
D = patches.shape[-1]
|
||
C = D // 4 # patch_size=2×2=4
|
||
x = rearrange(patches, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
||
h=h_len, w=w_len, c=C, ph=2, pw=2)
|
||
if extra_shape is not None:
|
||
B_orig, C_orig, T = extra_shape
|
||
# Unmerge B and T: [B_orig*T, C, H, W] → [B_orig, C, T, H, W]
|
||
H, W = x.shape[2], x.shape[3]
|
||
x = x.reshape(B_orig, T, C, H, W).permute(0, 2, 1, 3, 4)
|
||
return x
|