def register_attn_hooks(self, unet):
        """Registering hooks to obtain the attention maps for guidance"""

        attn_procs = {}

        for name in unet.attn_processors.keys():
            # Only obtain the queries and keys from cross-attention
            if name.endswith("attn1.processor") or name.endswith("fuser.attn.processor"):
                # Keep the same attn_processors for self-attention (no hooks for self-attention)
                attn_procs[name] = unet.attn_processors[name]
                continue

            cross_attention_dim = None if name.endswith("attn1.processor") else unet.config.cross_attention_dim

            if name.startswith("mid_block"):
                hidden_size = unet.config.block_out_channels[-1]
            elif name.startswith("up_blocks"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(unet.config.block_out_channels))[block_id]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = unet.config.block_out_channels[block_id]

            attn_procs[name] = AttnProcessorWithHook(
                attn_processor_key=name,
                hidden_size=hidden_size,
                cross_attention_dim=cross_attention_dim,
                hook=self.attn_hook,
                fast_attn=True,
                # Not enabled by default
                enabled=False,
            )

        unet.set_attn_processor(attn_procs)
Registering hooks to obtain the attention maps for guidance
examples/community/llm_grounded_diffusion.py