feat: add causal frame set attention fusion

infer_turn_head.py

@@ -30,6 +30,7 @@ from easycom_head import (
     output_space_for_loss,
     transition_targets_from_activity,
 )
+from local_turn_features import FramePersonFeatures
 
 
 def parse_args() -> argparse.Namespace:
@@ -118,6 +119,7 @@ def _predict_states(
     show_progress: bool,
     loss_name: str,
     local_features: torch.Tensor | None = None,
+    local_person_features: FramePersonFeatures | None = None,
 ) -> tuple[np.ndarray, np.ndarray | None]:
     chunks: list[torch.Tensor] = []
     transition_chunks: list[torch.Tensor] = []
@@ -139,11 +141,35 @@ def _predict_states(
             local_features[start : start + batch_size].to(device=device, dtype=torch.float32)
             if local_features is not None else None
         )
+        person_batch = None
+        person_mask = None
+        if local_person_features is not None:
+            rows = torch.arange(
+                start,
+                min(start + batch_size, len(states)),
+            )
+            person_batch, person_mask = local_person_features.padded(rows)
+            person_batch = person_batch.to(
+                device=device,
+                dtype=torch.float32,
+                non_blocking=True,
+            )
+            person_mask = person_mask.to(device=device, non_blocking=True)
         if head.transition_auxiliary:
-            activity_logits, transition_logits = head.forward_with_transition(batch, side)
+            activity_logits, transition_logits = head.forward_with_transition(
+                batch,
+                side,
+                person_batch,
+                person_mask,
+            )
             transition_chunks.append(transition_logits.sigmoid().cpu())
         else:
-            activity_logits = head(batch, side)
+            activity_logits = head(
+                batch,
+                side,
+                person_batch,
+                person_mask,
+            )
         chunks.append(
             marginal_probabilities_from_logits(activity_logits, loss_name).cpu()
         )
@@ -403,6 +429,7 @@ def main() -> None:
             local_feature_root=args.local_feature_root,
             local_modalities=local_contract["modalities"] if local_contract else "av",
             local_contract=local_contract,
+            local_fusion=head.local_fusion,
         )
         if data.hidden_size != int(head_config["hidden_size"]):
             raise ValueError(
@@ -423,6 +450,7 @@ def main() -> None:
             show_progress=not args.no_progress,
             loss_name=loss_name,
             local_features=data.local_features,
+            local_person_features=data.local_person_features,
         )
         predictions, valid = _align_to_egospeak_timeline(
             probabilities,
@@ -484,6 +512,7 @@ def main() -> None:
             "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
             "checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
             "layer_fusion": head.layer_fusion,
+            "local_fusion": head.local_fusion,
             "local_contract": local_contract,
             "transition_auxiliary": head.transition_auxiliary,
             "transition_objective": checkpoint_payload.get(