| 174 | |
| 175 | |
| 176 | class AttentionStore(AttentionControl): |
| 177 | |
| 178 | @staticmethod |
| 179 | def get_empty_store(): |
| 180 | return {"down_cross": [], "mid_cross": [], "up_cross": [], |
| 181 | "down_self": [], "mid_self": [], "up_self": []} |
| 182 | |
| 183 | def forward(self, attn, is_cross: bool, place_in_unet: str): |
| 184 | key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" |
| 185 | if attn.shape[1] <= 64 ** 2: # avoid memory overhead origin:32 ** 2 |
| 186 | self.step_store[key].append(attn) |
| 187 | return attn |
| 188 | |
| 189 | def between_steps(self): |
| 190 | if len(self.attention_store) == 0: |
| 191 | self.attention_store = self.step_store |
| 192 | if self.all_step_attention_store: |
| 193 | self.all_step_attention[self.cur_step-1] = {"down_cross": [], "mid_cross": [], "up_cross": [], |
| 194 | "down_self": [], "mid_self": [], "up_self": []} |
| 195 | for key in self.attention_store: |
| 196 | for i in range(len(self.attention_store[key])): |
| 197 | self.all_step_attention[self.cur_step-1][key].append(self.step_store[key][i].clone()) |
| 198 | else: |
| 199 | for key in self.attention_store: |
| 200 | for i in range(len(self.attention_store[key])): |
| 201 | self.attention_store[key][i] += self.step_store[key][i] |
| 202 | |
| 203 | if self.all_step_attention_store: |
| 204 | self.all_step_attention[self.cur_step-1] = self.step_store |
| 205 | |
| 206 | # |
| 207 | self.step_store = self.get_empty_store() |
| 208 | |
| 209 | |
| 210 | def get_average_attention(self): |
| 211 | average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store} |
| 212 | return average_attention |
| 213 | |
| 214 | |
| 215 | def reset(self): |
| 216 | super(AttentionStore, self).reset() |
| 217 | self.step_store = self.get_empty_store() |
| 218 | self.attention_store = {} |
| 219 | |
| 220 | def __init__(self, all_step_attention_store=False): |
| 221 | super(AttentionStore, self).__init__() |
| 222 | self.step_store = self.get_empty_store() |
| 223 | self.attention_store = {} |
| 224 | self.all_step_attention_store = all_step_attention_store |
| 225 | if all_step_attention_store: |
| 226 | self.all_step_attention = {} |
| 227 | |