text-rows-only full SDPA (exact)

app.py

@@ -68,6 +68,37 @@ FPS = 8  # the authors' output frame rate
 STEM = "scene"
 DEVICE = "cuda"
 
+# Exact inference-time speedup of the authors' attention processor: its full
+# SDPA's video rows are overwritten by the supervised decomposition, so only
+# the text rows of that SDPA are computed here (same output, ~half the attention).
+import torch.nn.functional as F  # noqa: E402
+from trajloc.backbones.cogvideox.attention import SupervisedJointAttnProcessor  # noqa: E402
+
+_orig_supervised_attention = SupervisedJointAttnProcessor._supervised_attention
+
+
+def _fast_supervised_attention(self, query, key, value, attention_mask, text_seq_length, batch_size):
+    if self._batched_mode or attention_mask is not None:
+        return _orig_supervised_attention(
+            self, query, key, value, attention_mask, text_seq_length, batch_size
+        )
+    text_out = F.scaled_dot_product_attention(
+        query[:, :, :text_seq_length], key, value, dropout_p=0.0, is_causal=False
+    )
+    video_out = self._supervised_video_rows(
+        query=query,
+        key=key,
+        value=value,
+        text_seq_length=text_seq_length,
+        supervised_positions=self.per_sample_positions[0],
+        heatmaps=self._flat_heatmaps.to(query.device, query.dtype).detach(),
+        batch_size=batch_size,
+    )
+    return torch.cat([text_out, video_out], dim=2)
+
+
+SupervisedJointAttnProcessor._supervised_attention = _fast_supervised_attention
+
 # --------------------------------------------------------------------- models
 ckpt_dir = Path(
     snapshot_download(