summaryrefslogtreecommitdiffstats
path: root/models/attention
diff options
context:
space:
mode:
authorVolpeon <git@volpeon.ink>2023-06-21 13:28:49 +0200
committerVolpeon <git@volpeon.ink>2023-06-21 13:28:49 +0200
commit8364ce697ddf6117fdd4f7222832d546d63880de (patch)
tree152c99815bbd8b2659d0dabe63c98f63151c97c2 /models/attention
parentFix LoRA training with DAdan (diff)
downloadtextual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.tar.gz
textual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.tar.bz2
textual-inversion-diff-8364ce697ddf6117fdd4f7222832d546d63880de.zip
Update
Diffstat (limited to 'models/attention')
-rw-r--r--models/attention/control.py106
-rw-r--r--models/attention/hook.py5
-rw-r--r--models/attention/structured.py65
3 files changed, 119 insertions, 57 deletions
diff --git a/models/attention/control.py b/models/attention/control.py
index 248bd9f..ec378c4 100644
--- a/models/attention/control.py
+++ b/models/attention/control.py
@@ -23,7 +23,7 @@ class AttentionControl(abc.ABC):
23 attn = self.forward(attn, is_cross, place_in_unet) 23 attn = self.forward(attn, is_cross, place_in_unet)
24 else: 24 else:
25 h = attn.shape[0] 25 h = attn.shape[0]
26 attn[h // 2:] = self.forward(attn[h // 2:], is_cross, place_in_unet) 26 attn[h // 2 :] = self.forward(attn[h // 2 :], is_cross, place_in_unet)
27 self.cur_att_layer += 1 27 self.cur_att_layer += 1
28 if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers: 28 if self.cur_att_layer == self.num_att_layers + self.num_uncond_att_layers:
29 self.cur_att_layer = 0 29 self.cur_att_layer = 0
@@ -49,12 +49,18 @@ class EmptyControl(AttentionControl):
49class AttentionStore(AttentionControl): 49class AttentionStore(AttentionControl):
50 @staticmethod 50 @staticmethod
51 def get_empty_store(): 51 def get_empty_store():
52 return {"down_cross": [], "mid_cross": [], "up_cross": [], 52 return {
53 "down_self": [], "mid_self": [], "up_self": []} 53 "down_cross": [],
54 "mid_cross": [],
55 "up_cross": [],
56 "down_self": [],
57 "mid_self": [],
58 "up_self": [],
59 }
54 60
55 def forward(self, attn, is_cross: bool, place_in_unet: str): 61 def forward(self, attn, is_cross: bool, place_in_unet: str):
56 key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" 62 key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
57 if attn.shape[1] <= 32 ** 2: # avoid memory overhead 63 if attn.shape[1] <= 32**2: # avoid memory overhead
58 self.step_store[key].append(attn) 64 self.step_store[key].append(attn)
59 return attn 65 return attn
60 66
@@ -68,8 +74,10 @@ class AttentionStore(AttentionControl):
68 self.step_store = self.get_empty_store() 74 self.step_store = self.get_empty_store()
69 75
70 def get_average_attention(self): 76 def get_average_attention(self):
71 average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] 77 average_attention = {
72 for key in self.attention_store} 78 key: [item / self.cur_step for item in self.attention_store[key]]
79 for key in self.attention_store
80 }
73 return average_attention 81 return average_attention
74 82
75 def reset(self): 83 def reset(self):
@@ -90,7 +98,7 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
90 return x_t 98 return x_t
91 99
92 def replace_self_attention(self, attn_base, att_replace): 100 def replace_self_attention(self, attn_base, att_replace):
93 if att_replace.shape[2] <= 16 ** 2: 101 if att_replace.shape[2] <= 16**2:
94 return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape) 102 return attn_base.unsqueeze(0).expand(att_replace.shape[0], *attn_base.shape)
95 else: 103 else:
96 return att_replace 104 return att_replace
@@ -101,41 +109,62 @@ class AttentionControlEdit(AttentionStore, abc.ABC):
101 109
102 def forward(self, attn, is_cross: bool, place_in_unet: str): 110 def forward(self, attn, is_cross: bool, place_in_unet: str):
103 super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet) 111 super(AttentionControlEdit, self).forward(attn, is_cross, place_in_unet)
104 if is_cross or (self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]): 112 if is_cross or (
113 self.num_self_replace[0] <= self.cur_step < self.num_self_replace[1]
114 ):
105 h = attn.shape[0] // (self.batch_size) 115 h = attn.shape[0] // (self.batch_size)
106 attn = attn.reshape(self.batch_size, h, *attn.shape[1:]) 116 attn = attn.reshape(self.batch_size, h, *attn.shape[1:])
107 attn_base, attn_repalce = attn[0], attn[1:] 117 attn_base, attn_repalce = attn[0], attn[1:]
108 if is_cross: 118 if is_cross:
109 alpha_words = self.cross_replace_alpha[self.cur_step] 119 alpha_words = self.cross_replace_alpha[self.cur_step]
110 attn_repalce_new = self.replace_cross_attention( 120 attn_repalce_new = (
111 attn_base, attn_repalce) * alpha_words + (1 - alpha_words) * attn_repalce 121 self.replace_cross_attention(attn_base, attn_repalce) * alpha_words
122 + (1 - alpha_words) * attn_repalce
123 )
112 attn[1:] = attn_repalce_new 124 attn[1:] = attn_repalce_new
113 else: 125 else:
114 attn[1:] = self.replace_self_attention(attn_base, attn_repalce) 126 attn[1:] = self.replace_self_attention(attn_base, attn_repalce)
115 attn = attn.reshape(self.batch_size * h, *attn.shape[2:]) 127 attn = attn.reshape(self.batch_size * h, *attn.shape[2:])
116 return attn 128 return attn
117 129
118 def __init__(self, prompts, num_steps: int, 130 def __init__(
119 cross_replace_steps: Union[float, Tuple[float, float], Dict[str, Tuple[float, float]]], 131 self,
120 self_replace_steps: Union[float, Tuple[float, float]], 132 prompts,
121 local_blend: Optional[LocalBlend]): 133 num_steps: int,
134 cross_replace_steps: Union[
135 float, Tuple[float, float], Dict[str, Tuple[float, float]]
136 ],
137 self_replace_steps: Union[float, Tuple[float, float]],
138 local_blend: Optional[LocalBlend],
139 ):
122 super(AttentionControlEdit, self).__init__() 140 super(AttentionControlEdit, self).__init__()
123 self.batch_size = len(prompts) 141 self.batch_size = len(prompts)
124 self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha( 142 self.cross_replace_alpha = ptp_utils.get_time_words_attention_alpha(
125 prompts, num_steps, cross_replace_steps, tokenizer).to(device) 143 prompts, num_steps, cross_replace_steps, tokenizer
144 ).to(device)
126 if type(self_replace_steps) is float: 145 if type(self_replace_steps) is float:
127 self_replace_steps = 0, self_replace_steps 146 self_replace_steps = 0, self_replace_steps
128 self.num_self_replace = int(num_steps * self_replace_steps[0]), int(num_steps * self_replace_steps[1]) 147 self.num_self_replace = int(num_steps * self_replace_steps[0]), int(
148 num_steps * self_replace_steps[1]
149 )
129 self.local_blend = local_blend 150 self.local_blend = local_blend
130 151
131 152
132class AttentionReplace(AttentionControlEdit): 153class AttentionReplace(AttentionControlEdit):
133 def replace_cross_attention(self, attn_base, att_replace): 154 def replace_cross_attention(self, attn_base, att_replace):
134 return torch.einsum('hpw,bwn->bhpn', attn_base, self.mapper) 155 return torch.einsum("hpw,bwn->bhpn", attn_base, self.mapper)
135 156
136 def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, 157 def __init__(
137 local_blend: Optional[LocalBlend] = None): 158 self,
138 super(AttentionReplace, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend) 159 prompts,
160 num_steps: int,
161 cross_replace_steps: float,
162 self_replace_steps: float,
163 local_blend: Optional[LocalBlend] = None,
164 ):
165 super(AttentionReplace, self).__init__(
166 prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
167 )
139 self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device) 168 self.mapper = seq_aligner.get_replacement_mapper(prompts, tokenizer).to(device)
140 169
141 170
@@ -145,9 +174,17 @@ class AttentionRefine(AttentionControlEdit):
145 attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas) 174 attn_replace = attn_base_replace * self.alphas + att_replace * (1 - self.alphas)
146 return attn_replace 175 return attn_replace
147 176
148 def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, 177 def __init__(
149 local_blend: Optional[LocalBlend] = None): 178 self,
150 super(AttentionRefine, self).__init__(prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend) 179 prompts,
180 num_steps: int,
181 cross_replace_steps: float,
182 self_replace_steps: float,
183 local_blend: Optional[LocalBlend] = None,
184 ):
185 super(AttentionRefine, self).__init__(
186 prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
187 )
151 self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer) 188 self.mapper, alphas = seq_aligner.get_refinement_mapper(prompts, tokenizer)
152 self.mapper, alphas = self.mapper.to(device), alphas.to(device) 189 self.mapper, alphas = self.mapper.to(device), alphas.to(device)
153 self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1]) 190 self.alphas = alphas.reshape(alphas.shape[0], 1, 1, alphas.shape[1])
@@ -156,13 +193,24 @@ class AttentionRefine(AttentionControlEdit):
156class AttentionReweight(AttentionControlEdit): 193class AttentionReweight(AttentionControlEdit):
157 def replace_cross_attention(self, attn_base, att_replace): 194 def replace_cross_attention(self, attn_base, att_replace):
158 if self.prev_controller is not None: 195 if self.prev_controller is not None:
159 attn_base = self.prev_controller.replace_cross_attention(attn_base, att_replace) 196 attn_base = self.prev_controller.replace_cross_attention(
197 attn_base, att_replace
198 )
160 attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :] 199 attn_replace = attn_base[None, :, :, :] * self.equalizer[:, None, None, :]
161 return attn_replace 200 return attn_replace
162 201
163 def __init__(self, prompts, num_steps: int, cross_replace_steps: float, self_replace_steps: float, equalizer, 202 def __init__(
164 local_blend: Optional[LocalBlend] = None, controller: Optional[AttentionControlEdit] = None): 203 self,
165 super(AttentionReweight, self).__init__(prompts, num_steps, 204 prompts,
166 cross_replace_steps, self_replace_steps, local_blend) 205 num_steps: int,
206 cross_replace_steps: float,
207 self_replace_steps: float,
208 equalizer,
209 local_blend: Optional[LocalBlend] = None,
210 controller: Optional[AttentionControlEdit] = None,
211 ):
212 super(AttentionReweight, self).__init__(
213 prompts, num_steps, cross_replace_steps, self_replace_steps, local_blend
214 )
167 self.equalizer = equalizer.to(device) 215 self.equalizer = equalizer.to(device)
168 self.prev_controller = controller 216 self.prev_controller = controller
diff --git a/models/attention/hook.py b/models/attention/hook.py
index 903de02..6b5fb68 100644
--- a/models/attention/hook.py
+++ b/models/attention/hook.py
@@ -3,6 +3,7 @@ import torch
3 3
4try: 4try:
5 import xformers.ops 5 import xformers.ops
6
6 xformers._is_functorch_available = True 7 xformers._is_functorch_available = True
7 MEM_EFFICIENT_ATTN = True 8 MEM_EFFICIENT_ATTN = True
8except ImportError: 9except ImportError:
@@ -42,10 +43,10 @@ def register_attention_control(model, controller):
42 return forward 43 return forward
43 44
44 def register_recr(net_, count, place_in_unet): 45 def register_recr(net_, count, place_in_unet):
45 if net_.__class__.__name__ == 'CrossAttention': 46 if net_.__class__.__name__ == "CrossAttention":
46 net_.forward = ca_forward(net_, place_in_unet) 47 net_.forward = ca_forward(net_, place_in_unet)
47 return count + 1 48 return count + 1
48 elif hasattr(net_, 'children'): 49 elif hasattr(net_, "children"):
49 for net__ in net_.children(): 50 for net__ in net_.children():
50 count = register_recr(net__, count, place_in_unet) 51 count = register_recr(net__, count, place_in_unet)
51 return count 52 return count
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):