diff --git a/comfy/ldm/cosmos/predict2.py b/comfy/ldm/cosmos/predict2.py index f560bded0..a8d4534d4 100644 --- a/comfy/ldm/cosmos/predict2.py +++ b/comfy/ldm/cosmos/predict2.py @@ -41,7 +41,7 @@ def torch_attention_op( q_B_S_H_D: torch.Tensor, k_B_S_H_D: torch.Tensor, v_B_S_H_D: torch.Tensor, - transformer_options: Optional[dict] = {}, + transformer_options: dict = {}, is_self_attention: bool = False, ) -> torch.Tensor: """Computes multi-head attention using PyTorch's native implementation. diff --git a/comfy/ldm/wan/ar_model.py b/comfy/ldm/wan/ar_model.py index 0d6fb313e..06c9da784 100644 --- a/comfy/ldm/wan/ar_model.py +++ b/comfy/ldm/wan/ar_model.py @@ -277,6 +277,7 @@ class CausalWanModel(WanModel): kv_caches=ar_state["kv_caches"], crossattn_caches=ar_state["crossattn_caches"], clip_fea=clip_fea, + transformer_options=transformer_options, ) return super().forward(x, timestep, context, clip_fea=clip_fea, diff --git a/tests-unit/comfy_test/attention_context_test.py b/tests-unit/comfy_test/attention_context_test.py new file mode 100644 index 000000000..df568eed4 --- /dev/null +++ b/tests-unit/comfy_test/attention_context_test.py @@ -0,0 +1,48 @@ +import torch + +from comfy.ldm.cosmos import predict2 +from comfy.ldm.wan.ar_model import CausalWanModel + + +def test_cosmos_attention_passes_self_attention_context(monkeypatch): + captured = {} + + def capture_attention(q, k, v, heads, **kwargs): + captured.update(kwargs) + return q + + monkeypatch.setattr(predict2, "optimized_attention", capture_attention) + q = torch.zeros((1, 2, 3, 4)) + + predict2.torch_attention_op(q, q, q, is_self_attention=True) + + assert captured["is_self_attention"] is True + assert captured["attention_token_shape"] == (2,) + + +def test_causal_wan_forward_passes_transformer_options_to_ar_block(): + transformer_options = { + "ar_state": { + "start_frame": 2, + "kv_caches": [object()], + "crossattn_caches": [object()], + }, + "optimized_attention_override": object(), + } + captured = {} + + class FakeCausalWan: + def forward_block(self, **kwargs): + captured.update(kwargs) + return "result" + + result = CausalWanModel.forward( + FakeCausalWan(), + x=torch.zeros((1, 4, 3, 2, 2)), + timestep=torch.zeros(1), + context=torch.zeros((1, 1, 1)), + transformer_options=transformer_options, + ) + + assert result == "result" + assert captured["transformer_options"] is transformer_options