diff --git a/headswap_ltx_masked_sampler.py b/headswap_ltx_masked_sampler.py index f0ec6de..91a6edb 100644 --- a/headswap_ltx_masked_sampler.py +++ b/headswap_ltx_masked_sampler.py @@ -325,6 +325,22 @@ def _mask_to_latent(masks, vae, latent_t, latent_h, latent_w): # node # ───────────────────────────────────────────────────────────────────────────── +def _keep_new_frames(samples, produced, total): + """Drop the head a chunk shares with the one before it. + + Chunks are sampled with an overlap for continuity, so consecutive chunks + cover some of the same frames. Concatenating them whole makes the clip + longer than it is and every frame after the first seam lands on the wrong + moment -- and, with per-frame crop boxes, in the wrong box. Keeping only + what each chunk adds makes the concatenation exactly `total` frames. + """ + have = samples.shape[2] + new = max(0, min(have, total - produced)) + if new < have: + samples = samples[:, :, have - new:] + return samples, produced + samples.shape[2] + + def _inject_transformer_options(guider, model_patcher, debug=False): """Copy the patched model's transformer_options INTO the guider's own dict. @@ -715,6 +731,7 @@ class BFSHeadSwapMaskedSampler: mc = LTXMultipleControls() out_latents = [] injections = [] + produced = 0 try: for idx, (a, b) in enumerate(chunks): g = guide[a:b] @@ -792,6 +809,7 @@ class BFSHeadSwapMaskedSampler: chunk = _Loop._sample_chunk(m, noise, sampler, sigmas, gd, latent, seed_offset=idx) got = chunk["samples"] if len(chunks) > 1 and not getattr(got, "is_nested", False): + got, produced = _keep_new_frames(got, produced, lat_t_total) # a long clip is many chunks: holding them all on the sampling # device grows VRAM with the clip while the model is still loaded got = got.to(comfy.model_management.intermediate_device()) diff --git a/pyproject.toml b/pyproject.toml index b622c08..2d41f4b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "bfsnodes" description = "Custom ComfyUI nodes for reference-guided video: the BFS Sampler (guide-driven IC-LoRAs with auto-configured crop, mask inpainting and temporal chunking in one node), LTX Multiple Controls and Identity Transfer (IPT2V reference + source-phase RoPE), LTXV Edit Anything, Head-Swap/Anime2Real Bernini conditioning, the AMV Guide Builder, and MiniMax-H3 single-frame VAE loaders with tile control." -version = "1.39.1" +version = "1.39.2" license = {file = "LICENSE"} dependencies = [ "torch", diff --git a/tests/test_headswap_auto_config.py b/tests/test_headswap_auto_config.py index 2d4a2d0..d9be262 100644 --- a/tests/test_headswap_auto_config.py +++ b/tests/test_headswap_auto_config.py @@ -249,6 +249,42 @@ class PasteBackWithoutCropTest(unittest.TestCase): self.assertEqual(out.shape, self.original.shape) +class ChunkOverlapTest(unittest.TestCase): + """Chunks are sampled with an overlap for continuity. Concatenating them + whole made the clip longer than it is, so every frame after the first seam + landed on the wrong moment -- and, with per-frame crop boxes, in the wrong + box. 97 frames at tile 72 / overlap 16 produced 15 latent frames where the + clip has 13.""" + + @staticmethod + def run_chunks(lat_per_chunk, total): + produced, out = 0, [] + for t in lat_per_chunk: + got, produced = NODE._keep_new_frames(torch.zeros(1, 4, t, 2, 2), produced, total) + out.append(got) + return torch.cat(out, dim=2) + + def test_the_concatenation_is_exactly_the_clip(self): + self.assertEqual(self.run_chunks([9, 6], 13).shape[2], 13) + + def test_it_keeps_each_chunk_s_tail_not_its_head(self): + a = torch.zeros(1, 4, 9, 2, 2) + b = torch.arange(6, dtype=torch.float32).view(1, 1, 6, 1, 1).expand(1, 4, 6, 2, 2).clone() + _, produced = NODE._keep_new_frames(a, 0, 13) + kept, _ = NODE._keep_new_frames(b, produced, 13) + self.assertEqual(kept.shape[2], 4) + self.assertEqual(kept[0, 0, 0, 0, 0].item(), 2.0) # frames 2..5, the new tail + + def test_a_single_chunk_is_untouched(self): + got, produced = NODE._keep_new_frames(torch.zeros(1, 4, 13, 2, 2), 0, 13) + self.assertEqual(got.shape[2], 13) + self.assertEqual(produced, 13) + + def test_it_never_returns_more_than_is_left(self): + got, _ = NODE._keep_new_frames(torch.zeros(1, 4, 9, 2, 2), 12, 13) + self.assertEqual(got.shape[2], 1) + + class NoShadowedModulesTest(unittest.TestCase): """`import comfy.x` inside a function binds `comfy` as a LOCAL for the whole function, so every earlier comfy.* use in it raises UnboundLocalError. It