[feature request] specify layer / timesteps

[enkeejunior1]

Thank you for your awesome work! It is easy to use without much modification.

BTW, since attention map often use as a tool for analyze the layer's behavior, specify what layer / timesteps to visualize greately improve capability.

[wooyeolBaek]

Thanks for the feature requests, and I hope my code has been helpful.

For layer specification, the module path and its attention map are saved as key and value in the attn_maps dictionary. By the way, it is a sum of all timesteps. For intuitive visualization, I simply summed the attention maps over all timesteps, just like in prompt-to-prompt. You can check it like this after the iteration ends:

# head_dim, seq_len, height, width
for k,v in attn_maps.items():
    print(k, v.shape)
    
> down_blocks.0.attentions.0.transformer_blocks.0.attn2 torch.Size([10, 77, 64, 96])
> ...

For timesteps, I’m thinking of adding an argument in the forward hook function to save attention maps at each timestep. I’ll notify you when it’s done.

[XLechter]

Thanks for the feature requests, and I hope my code has been helpful.

For layer specification, the module path and its attention map are saved as key and value in the attn_maps dictionary. By the way, it is a sum of all timesteps. For intuitive visualization, I simply summed the attention maps over all timesteps, just like in prompt-to-prompt. You can check it like this after the iteration ends:

# head_dim, seq_len, height, width
for k,v in attn_maps.items():
    print(k, v.shape)
    
> down_blocks.0.attentions.0.transformer_blocks.0.attn2 torch.Size([10, 77, 64, 96])
> ...

For timesteps, I’m thinking of adding an argument in the forward hook function to save attention maps at each timestep. I’ll notify you when it’s done.

@wooyeolBaek Thanks for sharing this. Looking forward to the attention map visualization by timesteps!

[wooyeolBaek]

@enkeejunior1 @XLechter
I've just added a feature to save attention maps for each timestep and path, and I've also updated the example. You can use it from now on.