Method __call__

src/attention_controllers.py:16–21

(self, attn, is_cross: bool = True, place_in_unet: str = None)

    # def forward(self, attn, is_cross: bool = True, place_in_unet: str = None):
    def __call__(self, attn, is_cross: bool = True, place_in_unet: str = None):
        # key = f"{place_in_unet}_{'cross' if is_cross else 'self'}"
        # if attn.shape[1] <= 32 ** 2:  # avoid memory overhead
        key = math.sqrt(attn.shape[1])
        self.step_store[key].append(attn)
        return