nodes/Danand-ComfyUI-ComfyCouple
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update attention_couple.py

Browse Source 
Ensure all dtypes match.
This commit is contained in:
Fox Rayside
2024-07-25 12:45:38 -04:00
committed by GitHub
parent 49e0628e73
commit 6dc99ad241
1 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
attention_couple.py
+5 -1
View File
@@ -140,7 +140,11 @@ class AttentionCouple:
q_target = q_list[i].repeat(length, 1, 1) q_target = q_list[i].repeat(length, 1, 1)
k = torch.cat([k[i].unsqueeze(0).repeat(b,1,1) for i in range(length)], dim=0) k = torch.cat([k[i].unsqueeze(0).repeat(b,1,1) for i in range(length)], dim=0)
v = torch.cat([v[i].unsqueeze(0).repeat(b,1,1) for i in range(length)], dim=0) v = torch.cat([v[i].unsqueeze(0).repeat(b,1,1) for i in range(length)], dim=0)
if k.dtype != q_target.dtype or v.dtype != q_target.dtype:
# Ensure all dtypes match
k = k.to(q_target.dtype)
v = v.to(q_target.dtype)
qkv = optimized_attention(q_target, k, v, extra_options["n_heads"]) qkv = optimized_attention(q_target, k, v, extra_options["n_heads"])
qkv = qkv * masks qkv = qkv * masks
qkv = qkv.view(length, b, -1, module.heads * module.dim_head).sum(dim=0) qkv = qkv.view(length, b, -1, module.heads * module.dim_head).sum(dim=0)