search

castorini / daam

Diffusion attentive attribution maps for interpreting Stable Diffusion.
MIT License
689 stars 63 forks source link

Trace hook never unhooks the attention processor #52

Closed rockerBOO closed 11 months ago

rockerBOO commented 12 months ago

https://github.com/castorini/daam/blob/a3231297bc3a6780a5ac7e56e3c5080badde4003/daam/trace.py#L281-L282

This hook doesn't unhook which can cause problems for the network in the future.

So I added a capture of the original processor and then added an unhook impl to set it back to the original processor. 

    def _hook_impl(self):
        self.original_processor = self.module.processor
        self.module.set_processor(self)

    def _unhook_impl(self):
        self.module.set_processor(self.original_processor)

I can make a PR in the future but in a messy situation at the moment.