diff options
Diffstat (limited to 'models/attention/control.py')
-rw-r--r-- | models/attention/control.py | 106 |
1 files changed, 77 insertions, 29 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): | |||
49 | class AttentionStore(AttentionControl): | 49 | class 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 | ||
132 | class AttentionReplace(AttentionControlEdit): | 153 | class 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): | |||
156 | class AttentionReweight(AttentionControlEdit): | 193 | class 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 |