fix: tile LDF decode conditions for conditional VAE

ldf_backend.py

@@ -187,7 +187,11 @@ class LDFVFIModel:
         target = torch.device(device)
         if target.type == "cuda" and not _cuda_bf16_supported(target):
             raise RuntimeError("LDF-VFI requires a CUDA GPU with BF16 support (Ampere or newer)")
-        self.transformer.to(device=target, dtype=self.dtype)
+        # from_pretrained(torch_dtype=...) keeps numerically sensitive modules
+        # (time embedding, norms, scale/shift) in FP32. Passing dtype here would
+        # flatten that mixed-precision policy and diffusers warns that results
+        # can become inconsistent.
+        self.transformer.to(device=target)
         self._move_auxiliary_models(target)
         self.device = str(target)
         return self
@@ -287,7 +291,46 @@ class LDFVFIModel:
         latent = rearrange(
             latent, "1 nt nh nw c t h w -> 1 nt c t (nh h) (nw w)"
         )
-        prediction = self.vae.decode(latent, dense, dense_mask)[..., :height, :width]
+        if dense.ndim != 5 or dense_mask.ndim != 5:
+            raise RuntimeError(
+                "LDF-VFI decode conditions must use [batch, channels, time, height, width]"
+            )
+
+        temporal_tiles = latent.shape[1]
+        if dense.shape[2] % temporal_tiles:
+            raise RuntimeError(
+                f"LDF-VFI condition length {dense.shape[2]} is not divisible by "
+                f"the {temporal_tiles} decode tiles"
+            )
+
+        # Mirror the official generate.vae_decode adapter. The conditional VAE
+        # consumes one condition and mask tile per latent temporal tile, not a
+        # single continuous 5-D condition tensor.
+        decode_height = latent.shape[-2] * self.vae.spatial_compression_ratio
+        decode_width = latent.shape[-1] * self.vae.spatial_compression_ratio
+        pad_height = decode_height - dense.shape[-2]
+        pad_width = decode_width - dense.shape[-1]
+        if pad_height < 0 or pad_width < 0:
+            raise RuntimeError(
+                "LDF-VFI decoded latent is smaller than its conditioning frames; "
+                "check the VAE tile and overlap settings"
+            )
+        dense = F.pad(dense, (0, pad_width, 0, pad_height))
+        dense = rearrange(
+            dense, "b c (nt t) h w -> b nt c t h w", nt=temporal_tiles
+        )
+
+        dense_mask = dense_mask[..., 0, 0]
+        dense_mask = repeat(
+            dense_mask, "b c t -> b c t h w", h=decode_height, w=decode_width
+        )
+        dense_mask = rearrange(
+            dense_mask, "b c (nt t) h w -> b nt c t h w", nt=temporal_tiles
+        )
+
+        prediction = self.vae.decode(
+            latent, dense, dense_mask
+        )[..., :height, :width]
         return rearrange(prediction, "1 c t h w -> t c h w").add(1).mul(0.5).clamp_(0, 1).float().cpu()
 
     @staticmethod

pyproject.toml

@@ -23,3 +23,7 @@ Repository = "https://github.com/Ethanfel/ComfyUI-Tween"
 [tool.comfy]
 PublisherId = "ethanfel"
 DisplayName = "Tween - Video Frame Interpolation"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "--confcutdir=tests"

tests/test_ldf_backend.py

@@ -0,0 +1,115 @@
+import torch
+import pytest
+
+from ldf_backend import LDFVFIModel
+
+
+class _RecordingConditionalVAE:
+    spatial_compression_ratio = 8
+
+    def __init__(self):
+        self.decode_shapes = None
+
+    def decode(self, latent, condition, mask):
+        self.decode_shapes = (
+            tuple(latent.shape),
+            tuple(condition.shape),
+            tuple(mask.shape),
+        )
+        # A small real result lets _decode finish without materializing the
+        # full 720p tensors used for the shape-only inputs above.
+        return torch.zeros(1, 3, 40, 1, 1)
+
+
+class _ShapeOnlyLDF(LDFVFIModel):
+    """Exercise sequence chunking without loading the multi-GB checkpoint."""
+
+    def __init__(self):
+        self.decode_tile_counts = []
+
+    def _prepare_condition(self, frames, mask, device):
+        assert frames.shape[0] == int(mask.sum())
+        dense = torch.empty(1, 3, self.TRAIN_FRAMES, 1, 1)
+        dense_mask = torch.empty(1, 1, self.TRAIN_FRAMES, 1, 1)
+        latent = torch.empty(
+            1, self.TRAIN_FRAMES // self.TILE_TIME, 1, 1, 1, 1, 1, 1
+        )
+        return dense, dense_mask, latent, latent
+
+    def _sample_free(self, condition, encoded_mask, schedule, device, progress):
+        return condition
+
+    def _sample_between(
+        self, previous, following, condition, encoded_mask,
+        schedule, t_cond, device, progress,
+    ):
+        return condition[:, self.CONDITION_TILES:-self.CONDITION_TILES]
+
+    def _sample_tail(
+        self, previous, condition, encoded_mask, schedule,
+        t_cond, device, progress,
+    ):
+        return condition[:, self.CONDITION_TILES:]
+
+    def _decode(self, latent, dense, dense_mask, height, width):
+        temporal_tiles = latent.shape[1]
+        assert dense.shape[2] == temporal_tiles * self.TILE_TIME
+        assert dense_mask.shape[2] == dense.shape[2]
+        self.decode_tile_counts.append(temporal_tiles)
+        return torch.empty(dense.shape[2], 3, height, width)
+
+
+def test_decode_tiles_40_frame_720p_condition_for_conditional_vae():
+    """Regression for the 5-D condition passed to upstream VAE.decode."""
+    model = LDFVFIModel.__new__(LDFVFIModel)
+    model.vae = _RecordingConditionalVAE()
+
+    # Representative first LDF output block at 720p. Meta tensors exercise
+    # exact shape transforms without allocating hundreds of MB in the test.
+    latent = torch.empty((1, 2, 4, 7, 4, 5, 24, 24), device="meta")
+    condition = torch.empty((1, 3, 40, 720, 1280), device="meta")
+    mask = torch.empty((1, 1, 40, 720, 1280), device="meta")
+
+    result = model._decode(latent, condition, mask, height=720, width=1280)
+
+    assert model.vae.decode_shapes == (
+        (1, 2, 4, 5, 96, 168),
+        (1, 2, 3, 20, 768, 1344),
+        (1, 2, 1, 20, 768, 1344),
+    )
+    assert result.shape == (40, 3, 1, 1)
+
+
+def test_decode_rejects_condition_length_that_cannot_tile():
+    model = LDFVFIModel.__new__(LDFVFIModel)
+    model.vae = _RecordingConditionalVAE()
+    latent = torch.empty((1, 2, 1, 1, 4, 5, 1, 1), device="meta")
+    condition = torch.empty((1, 3, 39, 8, 8), device="meta")
+    mask = torch.empty((1, 1, 39, 8, 8), device="meta")
+
+    try:
+        model._decode(latent, condition, mask, height=8, width=8)
+    except RuntimeError as error:
+        assert "not divisible" in str(error)
+    else:
+        raise AssertionError("Expected an invalid temporal tile error")
+
+
+@pytest.mark.parametrize("temporal_factor", [2, 3, 8, 16])
+def test_sequence_chunk_conditions_align_for_reported_391_frames(temporal_factor):
+    model = _ShapeOnlyLDF()
+    source = torch.empty(391, 3, 1, 1)
+    chunks = model._interpolate_sequence_impl(
+        source=source,
+        factor=temporal_factor,
+        schedule=torch.tensor([1.0, 0.0]),
+        t_cond=0.1,
+        device=torch.device("cpu"),
+        progress=lambda: None,
+    )
+
+    expected_frames = (source.shape[0] - 1) * temporal_factor + 1
+    assert sum(chunk.shape[0] for chunk in chunks) >= expected_frames
+    assert len(chunks) == model.sampling_block_count(
+        source.shape[0], temporal_factor
+    )