From b2350e7f081ceb0e6ed48066319c63e744121ac6 Mon Sep 17 00:00:00 2001 From: Ethan Fel <28683423+ethanfel@users.noreply.github.com> Date: Sat, 15 Aug 2026 21:46:17 +0200 Subject: [PATCH] fix: tile LDF decode conditions for conditional VAE --- ldf_backend.py | 47 +++++++++++++++- pyproject.toml | 4 ++ tests/test_ldf_backend.py | 115 ++++++++++++++++++++++++++++++++++++++ 3 files changed, 164 insertions(+), 2 deletions(-) create mode 100644 tests/test_ldf_backend.py diff --git a/ldf_backend.py b/ldf_backend.py index 51c2fe7..9e33e77 100644 --- a/ldf_backend.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 17bb1e9..06b6c64 100644 --- a/pyproject.toml +++ b/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" diff --git a/tests/test_ldf_backend.py b/tests/test_ldf_backend.py new file mode 100644 index 0000000..6f28fd3 --- /dev/null +++ b/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 + )