Fix i2v guide mask broadcasting

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")):

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