visualize of attention

[SunJ1025]

Hello bradyz, thanks for your great work! I'm doing some further research based on your work, I want to ask how to achieve the visualize of the cross attention? It really looks amazing.
Thanks in advance for your answer!