(self, attn, is_cross: bool, place_in_unet: str)
| 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: |
nothing calls this directly
no outgoing calls
no test coverage detected