diff options
author | Volpeon <git@volpeon.ink> | 2023-06-21 13:28:49 +0200 |
---|---|---|
committer | Volpeon <git@volpeon.ink> | 2023-06-21 13:28:49 +0200 |
commit | 8364ce697ddf6117fdd4f7222832d546d63880de (patch) | |
tree | 152c99815bbd8b2659d0dabe63c98f63151c97c2 /models/attention/structured.py | |
parent | Fix LoRA training with DAdan (diff) | |
download | textual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.tar.gz textual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.tar.bz2 textual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.zip |
Update
Diffstat (limited to 'models/attention/structured.py')
-rw-r--r-- | models/attention/structured.py | 65 |
1 files changed, 39 insertions, 26 deletions
diff --git a/models/attention/structured.py b/models/attention/structured.py index 24d889f..5bbbc06 100644 --- a/models/attention/structured.py +++ b/models/attention/structured.py | |||
@@ -16,7 +16,9 @@ class StructuredAttentionControl(AttentionControl): | |||
16 | if self.struct_attn: | 16 | if self.struct_attn: |
17 | out = self.struct_qkv(q, context, mask) | 17 | out = self.struct_qkv(q, context, mask) |
18 | else: | 18 | else: |
19 | context = torch.cat([context[0], context[1]['k'][0]], dim=0) # use key tensor for context | 19 | context = torch.cat( |
20 | [context[0], context[1]["k"][0]], dim=0 | ||
21 | ) # use key tensor for context | ||
20 | out = self.normal_qkv(q, context, mask) | 22 | out = self.normal_qkv(q, context, mask) |
21 | else: | 23 | else: |
22 | context = default(context, x) | 24 | context = default(context, x) |
@@ -29,11 +31,13 @@ class StructuredAttentionControl(AttentionControl): | |||
29 | context: list of [uc, list of conditional context] | 31 | context: list of [uc, list of conditional context] |
30 | """ | 32 | """ |
31 | uc_context = context[0] | 33 | uc_context = context[0] |
32 | context_k, context_v = context[1]['k'], context[1]['v'] | 34 | context_k, context_v = context[1]["k"], context[1]["v"] |
33 | 35 | ||
34 | if isinstance(context_k, list) and isinstance(context_v, list): | 36 | if isinstance(context_k, list) and isinstance(context_v, list): |
35 | out = self.multi_qkv(q, uc_context, context_k, context_v, mask) | 37 | out = self.multi_qkv(q, uc_context, context_k, context_v, mask) |
36 | elif isinstance(context_k, torch.Tensor) and isinstance(context_v, torch.Tensor): | 38 | elif isinstance(context_k, torch.Tensor) and isinstance( |
39 | context_v, torch.Tensor | ||
40 | ): | ||
37 | out = self.heterogeous_qkv(q, uc_context, context_k, context_v, mask) | 41 | out = self.heterogeous_qkv(q, uc_context, context_k, context_v, mask) |
38 | else: | 42 | else: |
39 | raise NotImplementedError | 43 | raise NotImplementedError |
@@ -50,36 +54,45 @@ class StructuredAttentionControl(AttentionControl): | |||
50 | k_c = [self.to_k(c_k) for c_k in context_k] | 54 | k_c = [self.to_k(c_k) for c_k in context_k] |
51 | v_c = [self.to_v(c_v) for c_v in context_v] | 55 | v_c = [self.to_v(c_v) for c_v in context_v] |
52 | 56 | ||
53 | q = rearrange(q, 'b n (h d) -> (b h) n d', h=h) | 57 | q = rearrange(q, "b n (h d) -> (b h) n d", h=h) |
54 | 58 | ||
55 | k_uc = rearrange(k_uc, 'b n (h d) -> (b h) n d', h=h) | 59 | k_uc = rearrange(k_uc, "b n (h d) -> (b h) n d", h=h) |
56 | v_uc = rearrange(v_uc, 'b n (h d) -> (b h) n d', h=h) | 60 | v_uc = rearrange(v_uc, "b n (h d) -> (b h) n d", h=h) |
57 | 61 | ||
58 | k_c = [rearrange(k, 'b n (h d) -> (b h) n d', h=h) for k in k_c] # NOTE: modification point | 62 | k_c = [ |
59 | v_c = [rearrange(v, 'b n (h d) -> (b h) n d', h=h) for v in v_c] | 63 | rearrange(k, "b n (h d) -> (b h) n d", h=h) for k in k_c |
64 | ] # NOTE: modification point | ||
65 | v_c = [rearrange(v, "b n (h d) -> (b h) n d", h=h) for v in v_c] | ||
60 | 66 | ||
61 | # get composition | 67 | # get composition |
62 | sim_uc = einsum('b i d, b j d -> b i j', q[:true_bs], k_uc) * self.scale | 68 | sim_uc = einsum("b i d, b j d -> b i j", q[:true_bs], k_uc) * self.scale |
63 | sim_c = [einsum('b i d, b j d -> b i j', q[true_bs:], k) * self.scale for k in k_c] | 69 | sim_c = [ |
70 | einsum("b i d, b j d -> b i j", q[true_bs:], k) * self.scale for k in k_c | ||
71 | ] | ||
64 | 72 | ||
65 | attn_uc = sim_uc.softmax(dim=-1) | 73 | attn_uc = sim_uc.softmax(dim=-1) |
66 | attn_c = [sim.softmax(dim=-1) for sim in sim_c] | 74 | attn_c = [sim.softmax(dim=-1) for sim in sim_c] |
67 | 75 | ||
68 | # get uc output | 76 | # get uc output |
69 | out_uc = einsum('b i j, b j d -> b i d', attn_uc, v_uc) | 77 | out_uc = einsum("b i j, b j d -> b i d", attn_uc, v_uc) |
70 | 78 | ||
71 | # get c output | 79 | # get c output |
72 | if len(v_c) == 1: | 80 | if len(v_c) == 1: |
73 | out_c_collect = [] | 81 | out_c_collect = [] |
74 | for attn in attn_c: | 82 | for attn in attn_c: |
75 | for v in v_c: | 83 | for v in v_c: |
76 | out_c_collect.append(einsum('b i j, b j d -> b i d', attn, v)) | 84 | out_c_collect.append(einsum("b i j, b j d -> b i d", attn, v)) |
77 | out_c = sum(out_c_collect) / len(out_c_collect) | 85 | out_c = sum(out_c_collect) / len(out_c_collect) |
78 | else: | 86 | else: |
79 | out_c = sum([einsum('b i j, b j d -> b i d', attn, v) for attn, v in zip(attn_c, v_c)]) / len(v_c) | 87 | out_c = sum( |
88 | [ | ||
89 | einsum("b i j, b j d -> b i d", attn, v) | ||
90 | for attn, v in zip(attn_c, v_c) | ||
91 | ] | ||
92 | ) / len(v_c) | ||
80 | 93 | ||
81 | out = torch.cat([out_uc, out_c], dim=0) | 94 | out = torch.cat([out_uc, out_c], dim=0) |
82 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) | 95 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h) |
83 | 96 | ||
84 | return out | 97 | return out |
85 | 98 | ||
@@ -88,21 +101,21 @@ class StructuredAttentionControl(AttentionControl): | |||
88 | k = self.to_k(context) | 101 | k = self.to_k(context) |
89 | v = self.to_v(context) | 102 | v = self.to_v(context) |
90 | 103 | ||
91 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) | 104 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) |
92 | 105 | ||
93 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale | 106 | sim = einsum("b i d, b j d -> b i j", q, k) * self.scale |
94 | 107 | ||
95 | if exists(mask): | 108 | if exists(mask): |
96 | mask = rearrange(mask, 'b ... -> b (...)') | 109 | mask = rearrange(mask, "b ... -> b (...)") |
97 | max_neg_value = -torch.finfo(sim.dtype).max | 110 | max_neg_value = -torch.finfo(sim.dtype).max |
98 | mask = repeat(mask, 'b j -> (b h) () j', h=h) | 111 | mask = repeat(mask, "b j -> (b h) () j", h=h) |
99 | sim.masked_fill_(~mask, max_neg_value) | 112 | sim.masked_fill_(~mask, max_neg_value) |
100 | 113 | ||
101 | # attention, what we cannot get enough of | 114 | # attention, what we cannot get enough of |
102 | attn = sim.softmax(dim=-1) | 115 | attn = sim.softmax(dim=-1) |
103 | 116 | ||
104 | out = einsum('b i j, b j d -> b i d', attn, v) | 117 | out = einsum("b i j, b j d -> b i d", attn, v) |
105 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) | 118 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h) |
106 | 119 | ||
107 | return out | 120 | return out |
108 | 121 | ||
@@ -111,21 +124,21 @@ class StructuredAttentionControl(AttentionControl): | |||
111 | k = self.to_k(torch.cat([uc_context, context_k], dim=0)) | 124 | k = self.to_k(torch.cat([uc_context, context_k], dim=0)) |
112 | v = self.to_v(torch.cat([uc_context, context_v], dim=0)) | 125 | v = self.to_v(torch.cat([uc_context, context_v], dim=0)) |
113 | 126 | ||
114 | q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v)) | 127 | q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)) |
115 | 128 | ||
116 | sim = einsum('b i d, b j d -> b i j', q, k) * self.scale | 129 | sim = einsum("b i d, b j d -> b i j", q, k) * self.scale |
117 | 130 | ||
118 | if exists(mask): | 131 | if exists(mask): |
119 | mask = rearrange(mask, 'b ... -> b (...)') | 132 | mask = rearrange(mask, "b ... -> b (...)") |
120 | max_neg_value = -torch.finfo(sim.dtype).max | 133 | max_neg_value = -torch.finfo(sim.dtype).max |
121 | mask = repeat(mask, 'b j -> (b h) () j', h=h) | 134 | mask = repeat(mask, "b j -> (b h) () j", h=h) |
122 | sim.masked_fill_(~mask, max_neg_value) | 135 | sim.masked_fill_(~mask, max_neg_value) |
123 | 136 | ||
124 | # attention, what we cannot get enough of | 137 | # attention, what we cannot get enough of |
125 | attn = sim.softmax(dim=-1) | 138 | attn = sim.softmax(dim=-1) |
126 | 139 | ||
127 | out = einsum('b i j, b j d -> b i d', attn, v) | 140 | out = einsum("b i j, b j d -> b i d", attn, v) |
128 | out = rearrange(out, '(b h) n d -> b n (h d)', h=h) | 141 | out = rearrange(out, "(b h) n d -> b n (h d)", h=h) |
129 | return out | 142 | return out |
130 | 143 | ||
131 | def get_kv(self, context): | 144 | def get_kv(self, context): |