| 5 | |
| 6 | |
| 7 | class CDAttentionStore(abc.ABC): |
| 8 | |
| 9 | @staticmethod |
| 10 | def get_empty_store(): |
| 11 | return {8: [], 16: [], 32: [], 64: []} |
| 12 | # return {"down_cross": [], "mid_cross": [], "up_cross": [], |
| 13 | # "down_self": [], "mid_self": [], "up_self": []} |
| 14 | |
| 15 | # def forward(self, attn, is_cross: bool = True, place_in_unet: str = None): |
| 16 | def __call__(self, attn, is_cross: bool = True, place_in_unet: str = None): |
| 17 | # key = f"{place_in_unet}_{'cross' if is_cross else 'self'}" |
| 18 | # if attn.shape[1] <= 32 ** 2: # avoid memory overhead |
| 19 | key = math.sqrt(attn.shape[1]) |
| 20 | self.step_store[key].append(attn) |
| 21 | return |
| 22 | |
| 23 | |
| 24 | |
| 25 | def between_steps(self): |
| 26 | # if len(self.attention_store) == 0: |
| 27 | # self.attention_store = self.step_store |
| 28 | # else: |
| 29 | # for key in self.attention_store: |
| 30 | # for i in range(len(self.attention_store[key])): |
| 31 | # self.attention_store[key][i] += self.step_store[key][i] |
| 32 | self.step_store = self.get_empty_store() |
| 33 | |
| 34 | # def get_average_attention(self): |
| 35 | # # average_attention = {key: [item / self.cur_step for item in self.attention_store[key]] for key in self.attention_store} |
| 36 | # average_attention = {key: [item for item in self.attention_store[key]] for key in self.attention_store} |
| 37 | |
| 38 | # return average_attention |
| 39 | |
| 40 | |
| 41 | def reset(self): |
| 42 | # super(CDAttentionStore, self).reset() |
| 43 | self.step_store = self.get_empty_store() |
| 44 | # self.attention_store = {} |
| 45 | |
| 46 | def __init__(self): |
| 47 | # super(CDAttentionStore, self).__init__() |
| 48 | self.step_store = self.get_empty_store() |
| 49 | # self.attention_store = {} |
| 50 | # ############################################################################################################# |
| 51 | # class AttentionControl(abc.ABC): |
| 52 | |