diff --git a/core/patchify.py b/core/patchify.py index 8ceaf6e..bc3cbd5 100644 --- a/core/patchify.py +++ b/core/patchify.py @@ -1,6 +1,10 @@ -"""Patchify/unpatchify for FLUX-family latent tensors (patch_size=2, auto-detect channels).""" +"""Patchify/unpatchify for latent tensors (patch_size=2, auto-detect channels). + +Handles 3D, 4D, and 5D inputs. Pads odd spatial dims to even before patchifying. +""" from einops import rearrange +import torch.nn.functional as F def patchify(x): @@ -11,25 +15,44 @@ def patchify(x): - 4D [B, C, H, W]: standard path, extra_shape=None - 5D [B, C, T, H, W]: video VAE, merges T into batch, extra_shape=(B, C, T) - L = (H/2) * (W/2), d = C * 2 * 2. + Pads odd H/W to even before patchifying. The pad amounts are stored + in the returned extra_shape for unpatchify to crop back. + + L = (H_padded/2) * (W_padded/2), d = C * 2 * 2. """ extra_shape = None + pad_h = 0 + pad_w = 0 + if x.ndim == 3: - # No batch dimension (e.g. LTXAV): [C, H, W] → [1, C, H, W] extra_shape = "unbatched" x = x.unsqueeze(0) elif 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 - if H < 2 or W < 2: - # Incompatible latent format (e.g. LTXAV uses flattened 1D layout) + if H < 1 or W < 1: return None, None, None, None - h_len = H // 2 - w_len = W // 2 + + # Pad odd dimensions to even (replicate last row/col) + if H % 2 != 0: + pad_h = 1 + if W % 2 != 0: + pad_w = 1 + if pad_h or pad_w: + x = F.pad(x, (0, pad_w, 0, pad_h), mode="replicate") + + H_p, W_p = x.shape[2], x.shape[3] + h_len = H_p // 2 + w_len = W_p // 2 patches = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=2, pw=2) + + # Bundle pad info with extra_shape + if pad_h or pad_w: + extra_shape = {"orig_extra": extra_shape, "pad_h": pad_h, "pad_w": pad_w} + return patches, h_len, w_len, extra_shape @@ -37,18 +60,34 @@ def unpatchify(patches, h_len, w_len, extra_shape=None): """Convert patch sequence [B, L, C*4] → latent, restoring original shape. Auto-detects channel count from patch dimension: C = D / 4. - Restores 3D/5D format based on extra_shape from patchify. + Handles padding removal and 3D/5D restoration based on extra_shape. """ 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 == "unbatched": - # Restore [1, C, H, W] → [C, H, W] + + # Unwrap pad info if present + pad_h = 0 + pad_w = 0 + orig_extra = extra_shape + if isinstance(extra_shape, dict): + pad_h = extra_shape["pad_h"] + pad_w = extra_shape["pad_w"] + orig_extra = extra_shape["orig_extra"] + + # Remove padding + if pad_h: + x = x[:, :, :-pad_h, :] + if pad_w: + x = x[:, :, :, :-pad_w] + + # Restore original format + if orig_extra == "unbatched": x = x.squeeze(0) - elif 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] + elif orig_extra is not None: + B_orig, C_orig, T = orig_extra H, W = x.shape[2], x.shape[3] - x = x.reshape(B_orig, T, C, H, W).permute(0, 2, 1, 3, 4) + x = x.reshape(B_orig, T, C_orig, H, W).permute(0, 2, 1, 3, 4) + return x