Pad odd spatial dims in patchify for LTXV compatibility

LTXV can produce latents with odd H or W (e.g. 11x20) which can't be
divided into 2x2 patches. Patchify now pads odd dims to even using
replicate padding before rearranging, and unpatchify crops back to the
original size. This keeps patch_size=2 consistent between calibration
and inference for all models.

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