diff --git a/samplers.py b/samplers.py index deab6678..09d39c2c 100644 --- a/samplers.py +++ b/samplers.py @@ -901,6 +901,22 @@ class FunPackLTXAVSceneChainSampler: tensor[tuple(slices)] = value return tensor + def _expand_mask_like(self, mask, target): + if mask.shape[0] != target.shape[0] or mask.shape[2] != target.shape[2]: + raise ValueError("Guide mask batch/time dimensions must match target mask.") + shape = list(mask.shape) + while len(shape) < target.dim(): + shape.append(1) + mask = mask.reshape(shape) + expand_shape = list(target.shape) + for dim in range(target.dim()): + if dim == 2: + continue + if mask.shape[dim] not in (1, target.shape[dim]): + raise ValueError("Guide mask dimensions are not broadcastable to target mask.") + expand_shape[dim] = target.shape[dim] + return mask.expand(expand_shape) + def _make_mask_tensor(self, tensor, overlap): mask = torch.ones_like(tensor) if overlap > 0: @@ -1011,8 +1027,8 @@ class FunPackLTXAVSceneChainSampler: if out_masks[0] is None: out_masks[0] = torch.ones_like(out_tensors[0]) out_tensors[0] = torch.cat([out_tensors[0], guide], dim=2) - if guide_mask.shape[1] != out_masks[0].shape[1]: - guide_mask = guide_mask.expand(-1, out_masks[0].shape[1], -1, -1, -1) + target_mask = self._time_slice(out_masks[0], 0, protected).to(guide_mask.device, guide_mask.dtype) + guide_mask = self._expand_mask_like(guide_mask, target_mask) out_masks[0] = torch.cat([out_masks[0].to(guide_mask.device, guide_mask.dtype), guide_mask], dim=2) if self._is_nested(chunk.get("samples")): diff --git a/tests/test_scene_chain_sampler.py b/tests/test_scene_chain_sampler.py index 46599629..52c5c1a4 100644 --- a/tests/test_scene_chain_sampler.py +++ b/tests/test_scene_chain_sampler.py @@ -257,6 +257,39 @@ def test_scene_chain_can_append_i2v_template_as_hidden_guide(): assert "i2v guide tokens=1 latent frame(s)" in status +def test_scene_chain_expands_compact_i2v_guide_mask_to_spatial_chunk_mask(): + sample_calls.clear() + node = FunPackLTXAVSceneChainSampler() + samples = torch.zeros(1, 2, 5, 24, 3) + samples[:, :, 0] = 7.0 + mask = torch.ones(1, 1, 5, 1, 1) + mask[:, :, 0] = 0.0 + latent_template = {"samples": samples, "noise_mask": mask} + positive = [scene_cond(0), scene_cond(1)] + + _, status, scene_count, _ = node.sample( + model=object(), + vae=FakeVAE(), + positive=positive, + negative=[], + sampler=object(), + sigmas=torch.tensor([1.0, 0.0]), + seed=45, + latent_template=latent_template, + num_frames_per_scene=5, + frame_overlap=2, + cfg=1.0, + max_scenes=2, + carry_i2v_guides=True, + ) + + second_call = sample_calls[1] + assert scene_count == 2 + assert second_call["noise_mask"].shape == second_call["latent_image"].shape + assert torch.all(second_call["noise_mask"][:, :, 5] == 0.0) + assert "i2v guide tokens=1 latent frame(s)" in status + + def test_scene_chain_does_not_carry_i2v_guides_by_default(): inputs = FunPackLTXAVSceneChainSampler.INPUT_TYPES()["required"]["carry_i2v_guides"][1] assert inputs["default"] is False