(self)
| 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): |
no test coverage detected