mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-18 20:38:15 +08:00
49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
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
|