lllyasviel/stable-diffusion-webui-forge
1
0
Fork 0
mirror of https://github.com/lllyasviel/stable-diffusion-webui-forge.git synced 2026-02-19 06:19:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
layerdiffusion
2024-08-08 16:19:17 -07:00
parent 8d336c3247
commit 26b7fea8a1
1 changed files with 3 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
backend/nn/flux.py
View File

@@ -129,8 +129,7 @@ class SelfAttention(nn.Module):
# q, k, v = rearrange(qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
B, L, _ = qkv.shape
qkv = qkv.view(B, L, 3, self.num_heads, -1) # Split into Q, K, V
qkv = qkv.permute(2, 0, 3, 1, 4) # Rearrange to (K B H L D)
q, k, v = qkv[0], qkv[1], qkv[2] # Separate Q, K, V
q, k, v = qkv.permute(2, 0, 3, 1, 4) # Rearrange to (K B H L D)
del qkv
q, k = self.norm(q, k, v)
@@ -196,10 +195,7 @@ class DoubleStreamBlock(nn.Module):
B, L, _ = img_qkv.shape
H = self.num_heads
D = img_qkv.shape[-1] // (3 * H)
img_qkv = img_qkv.view(B, L, 3, H, D)
img_q = img_qkv[:, :, 0, :, :].permute(2, 0, 1, 3)
img_k = img_qkv[:, :, 1, :, :].permute(2, 0, 1, 3)
img_v = img_qkv[:, :, 2, :, :].permute(2, 0, 1, 3)
img_q, img_k, img_v = img_qkv.view(B, L, 3, H, D).permute(2, 0, 1, 3)
del img_qkv
img_q, img_k = self.img_attn.norm(img_q, img_k, img_v)
@@ -212,10 +208,7 @@ class DoubleStreamBlock(nn.Module):
# txt_q, txt_k, txt_v = rearrange(txt_qkv, "B L (K H D) -> K B H L D", K=3, H=self.num_heads)
B, L, _ = txt_qkv.shape
txt_qkv = txt_qkv.view(B, L, 3, H, D)
txt_q = txt_qkv[:, :, 0, :, :].permute(2, 0, 1, 3)
txt_k = txt_qkv[:, :, 1, :, :].permute(2, 0, 1, 3)
txt_v = txt_qkv[:, :, 2, :, :].permute(2, 0, 1, 3)
txt_q, txt_k, txt_v = txt_qkv.view(B, L, 3, H, D).permute(2, 0, 1, 3)
del txt_qkv
txt_q, txt_k = self.txt_attn.norm(txt_q, txt_k, txt_v)