utils/attention_utils.py

import numpy as np
import os
import matplotlib as mpl
import matplotlib.pyplot as plt
import seaborn as sns
import torch
import torchvision

from utils.richtext_utils import seed_everything
from sklearn.cluster import KMeans, SpectralClustering

# SelfAttentionLayers = [
#     # 'down_blocks.0.attentions.0.transformer_blocks.0.attn1',
#     # 'down_blocks.0.attentions.1.transformer_blocks.0.attn1',
#     'down_blocks.1.attentions.0.transformer_blocks.0.attn1',
#     # 'down_blocks.1.attentions.1.transformer_blocks.0.attn1',
#     'down_blocks.2.attentions.0.transformer_blocks.0.attn1',
#     'down_blocks.2.attentions.1.transformer_blocks.0.attn1',
#     'mid_block.attentions.0.transformer_blocks.0.attn1',
#     'up_blocks.1.attentions.0.transformer_blocks.0.attn1',
#     'up_blocks.1.attentions.1.transformer_blocks.0.attn1',
#     'up_blocks.1.attentions.2.transformer_blocks.0.attn1',
#     # 'up_blocks.2.attentions.0.transformer_blocks.0.attn1',
#     'up_blocks.2.attentions.1.transformer_blocks.0.attn1',
#     # 'up_blocks.2.attentions.2.transformer_blocks.0.attn1',
#     # 'up_blocks.3.attentions.0.transformer_blocks.0.attn1',
#     # 'up_blocks.3.attentions.1.transformer_blocks.0.attn1',
#     # 'up_blocks.3.attentions.2.transformer_blocks.0.attn1',
# ]

SelfAttentionLayers = [
    # 'down_blocks.0.attentions.0.transformer_blocks.0.attn1',
    # 'down_blocks.0.attentions.1.transformer_blocks.0.attn1',
    'down_blocks.1.attentions.0.transformer_blocks.0.attn1',
    # 'down_blocks.1.attentions.1.transformer_blocks.0.attn1',
    'down_blocks.2.attentions.0.transformer_blocks.0.attn1',
    'down_blocks.2.attentions.1.transformer_blocks.0.attn1',
    'mid_block.attentions.0.transformer_blocks.0.attn1',
    'up_blocks.1.attentions.0.transformer_blocks.0.attn1',
    'up_blocks.1.attentions.1.transformer_blocks.0.attn1',
    'up_blocks.1.attentions.2.transformer_blocks.0.attn1',
    # 'up_blocks.2.attentions.0.transformer_blocks.0.attn1',
    'up_blocks.2.attentions.1.transformer_blocks.0.attn1',
    # 'up_blocks.2.attentions.2.transformer_blocks.0.attn1',
    # 'up_blocks.3.attentions.0.transformer_blocks.0.attn1',
    # 'up_blocks.3.attentions.1.transformer_blocks.0.attn1',
    # 'up_blocks.3.attentions.2.transformer_blocks.0.attn1',
]


CrossAttentionLayers = [
    # 'down_blocks.0.attentions.0.transformer_blocks.0.attn2',
    # 'down_blocks.0.attentions.1.transformer_blocks.0.attn2',
    'down_blocks.1.attentions.0.transformer_blocks.0.attn2',
    # 'down_blocks.1.attentions.1.transformer_blocks.0.attn2',
    'down_blocks.2.attentions.0.transformer_blocks.0.attn2',
    'down_blocks.2.attentions.1.transformer_blocks.0.attn2',
    'mid_block.attentions.0.transformer_blocks.0.attn2',
    'up_blocks.1.attentions.0.transformer_blocks.0.attn2',
    'up_blocks.1.attentions.1.transformer_blocks.0.attn2',
    'up_blocks.1.attentions.2.transformer_blocks.0.attn2',
    # 'up_blocks.2.attentions.0.transformer_blocks.0.attn2',
    'up_blocks.2.attentions.1.transformer_blocks.0.attn2',
    # 'up_blocks.2.attentions.2.transformer_blocks.0.attn2',
    # 'up_blocks.3.attentions.0.transformer_blocks.0.attn2',
    # 'up_blocks.3.attentions.1.transformer_blocks.0.attn2',
    # 'up_blocks.3.attentions.2.transformer_blocks.0.attn2'
]

# CrossAttentionLayers = [
#     'down_blocks.0.attentions.0.transformer_blocks.0.attn2',
#     'down_blocks.0.attentions.1.transformer_blocks.0.attn2',
#     'down_blocks.1.attentions.0.transformer_blocks.0.attn2',
#     'down_blocks.1.attentions.1.transformer_blocks.0.attn2',
#     'down_blocks.2.attentions.0.transformer_blocks.0.attn2',
#     'down_blocks.2.attentions.1.transformer_blocks.0.attn2',
#     'mid_block.attentions.0.transformer_blocks.0.attn2',
#     'up_blocks.1.attentions.0.transformer_blocks.0.attn2',
#     'up_blocks.1.attentions.1.transformer_blocks.0.attn2',
#     'up_blocks.1.attentions.2.transformer_blocks.0.attn2',
#     'up_blocks.2.attentions.0.transformer_blocks.0.attn2',
#     'up_blocks.2.attentions.1.transformer_blocks.0.attn2',
#     'up_blocks.2.attentions.2.transformer_blocks.0.attn2',
#     'up_blocks.3.attentions.0.transformer_blocks.0.attn2',
#     'up_blocks.3.attentions.1.transformer_blocks.0.attn2',
#     'up_blocks.3.attentions.2.transformer_blocks.0.attn2'
# ]

# CrossAttentionLayers_XL = [
#     'up_blocks.0.attentions.0.transformer_blocks.1.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.2.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.3.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.4.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.5.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.6.attn2',
#     'up_blocks.0.attentions.0.transformer_blocks.7.attn2',
# ]
CrossAttentionLayers_XL = [
    'down_blocks.2.attentions.1.transformer_blocks.3.attn2',
    'down_blocks.2.attentions.1.transformer_blocks.4.attn2',
    'mid_block.attentions.0.transformer_blocks.0.attn2',
    'mid_block.attentions.0.transformer_blocks.1.attn2',
    'mid_block.attentions.0.transformer_blocks.2.attn2',
    'mid_block.attentions.0.transformer_blocks.3.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.1.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.2.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.3.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.4.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.5.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.6.attn2',
    'up_blocks.0.attentions.0.transformer_blocks.7.attn2',
    'up_blocks.1.attentions.0.transformer_blocks.0.attn2'
]

def split_attention_maps_over_steps(attention_maps):
    r"""Function for splitting attention maps over steps.
    Args:
        attention_maps (dict): Dictionary of attention maps.
        sampler_order (int): Order of the sampler.
    """
    # This function splits attention maps into unconditional and conditional score and over steps

    attention_maps_cond = dict()    # Maps corresponding to conditional score
    attention_maps_uncond = dict()  # Maps corresponding to unconditional score

    for layer in attention_maps.keys():

        for step_num in range(len(attention_maps[layer])):
            if step_num not in attention_maps_cond:
                attention_maps_cond[step_num] = dict()
                attention_maps_uncond[step_num] = dict()

            attention_maps_uncond[step_num].update(
                {layer: attention_maps[layer][step_num][:1]})
            attention_maps_cond[step_num].update(
                {layer: attention_maps[layer][step_num][1:2]})

    return attention_maps_cond, attention_maps_uncond


def save_attention_heatmaps(attention_maps, tokens_vis, save_dir, prefix):
    r"""Function to plot heatmaps for attention maps.

    Args:
        attention_maps (dict): Dictionary of attention maps per layer
        save_dir (str): Directory to save attention maps
        prefix (str): Filename prefix for html files

    Returns:
        Heatmaps, one per sample.
    """

    html_names = []

    idx = 0
    html_list = []

    for layer in attention_maps.keys():
        if idx == 0:
            # import ipdb;ipdb.set_trace()
            # create a set of html files.

            batch_size = attention_maps[layer].shape[0]

            for sample_num in range(batch_size):
                # html path
                html_rel_path = os.path.join('sample_{}'.format(
                    sample_num), '{}.html'.format(prefix))
                html_names.append(html_rel_path)
                html_path = os.path.join(save_dir, html_rel_path)
                os.makedirs(os.path.dirname(html_path), exist_ok=True)
                html_list.append(open(html_path, 'wt'))
                html_list[sample_num].write(
                    '<html><head></head><body><table>\n')

        for sample_num in range(batch_size):

            save_path = os.path.join(save_dir, 'sample_{}'.format(sample_num),
                                     prefix, 'layer_{}'.format(layer)) + '.jpg'
            Path(os.path.dirname(save_path)).mkdir(parents=True, exist_ok=True)

            layer_name = 'layer_{}'.format(layer)
            html_list[sample_num].write(
                f'<tr><td><h1>{layer_name}</h1></td></tr>\n')

            prefix_stem = prefix.split('/')[-1]
            relative_image_path = os.path.join(
                prefix_stem, 'layer_{}'.format(layer)) + '.jpg'
            html_list[sample_num].write(
                f'<tr><td><img src=\"{relative_image_path}\"></td></tr>\n')

            plt.figure()
            plt.clf()
            nrows = 2
            ncols = 7
            fig, axs = plt.subplots(nrows=nrows, ncols=ncols)

            fig.set_figheight(8)
            fig.set_figwidth(28.5)

            # axs[0].set_aspect('equal')
            # axs[1].set_aspect('equal')
            # axs[2].set_aspect('equal')
            # axs[3].set_aspect('equal')
            # axs[4].set_aspect('equal')
            # axs[5].set_aspect('equal')

            cmap = plt.get_cmap('YlOrRd')

            for rid in range(nrows):
                for cid in range(ncols):
                    tid = rid*ncols + cid
                    # import ipdb;ipdb.set_trace()
                    attention_map_cur = attention_maps[layer][sample_num, :, :, tid].numpy(
                    )
                    vmax = float(attention_map_cur.max())
                    vmin = float(attention_map_cur.min())
                    sns.heatmap(
                        attention_map_cur, annot=False, cbar=False, ax=axs[rid, cid],
                        cmap=cmap, vmin=vmin, vmax=vmax
                    )
                    axs[rid, cid].set_xlabel(tokens_vis[tid])

            # axs[0].set_xlabel('Self attention')
            # axs[1].set_xlabel('Temporal attention')
            # axs[2].set_xlabel('T5 text attention')
            # axs[3].set_xlabel('CLIP text attention')
            # axs[4].set_xlabel('CLIP image attention')
            # axs[5].set_xlabel('Null text token')

            norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
            sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
            # fig.colorbar(sm, cax=axs[6])

            fig.tight_layout()
            plt.savefig(save_path, dpi=64)
            plt.close('all')

        if idx == (len(attention_maps.keys()) - 1):
            for sample_num in range(batch_size):
                html_list[sample_num].write('</table></body></html>')
                html_list[sample_num].close()

        idx += 1

    return html_names


def create_recursive_html_link(html_path, save_dir):
    r"""Function for creating recursive html links.
    If the path is dir1/dir2/dir3/*.html,
    we create chained directories
        -dir1
            dir1.html (has links to all children)
            -dir2
                dir2.html   (has links to all children)
                -dir3
                    dir3.html

    Args:
        html_path (str): Path to html file.
        save_dir (str): Save directory.
    """

    html_path_split = os.path.splitext(html_path)[0].split('/')
    if len(html_path_split) == 1:
        return

    # First create the root directory
    root_dir = html_path_split[0]
    child_dir = html_path_split[1]

    cur_html_path = os.path.join(save_dir, '{}.html'.format(root_dir))
    if os.path.exists(cur_html_path):

        fp = open(cur_html_path, 'r')
        lines_written = fp.readlines()
        fp.close()

        fp = open(cur_html_path, 'a+')
        child_path = os.path.join(root_dir, f'{child_dir}.html')
        line_to_write = f'<tr><td><a href=\"{child_path}\">{child_dir}</a></td></tr>\n'

        if line_to_write not in lines_written:
            fp.write('<html><head></head><body><table>\n')
            fp.write(line_to_write)
            fp.write('</table></body></html>')
        fp.close()

    else:

        fp = open(cur_html_path, 'w')

        child_path = os.path.join(root_dir, f'{child_dir}.html')
        line_to_write = f'<tr><td><a href=\"{child_path}\">{child_dir}</a></td></tr>\n'

        fp.write('<html><head></head><body><table>\n')
        fp.write(line_to_write)
        fp.write('</table></body></html>')

        fp.close()

    child_path = '/'.join(html_path.split('/')[1:])
    save_dir = os.path.join(save_dir, root_dir)
    create_recursive_html_link(child_path, save_dir)


def visualize_attention_maps(attention_maps_all, save_dir, width, height, tokens_vis):
    r"""Function to visualize attention maps.
    Args:
        save_dir (str): Path to save attention maps
        batch_size (int): Batch size
        sampler_order (int): Sampler order
    """

    rand_name = list(attention_maps_all.keys())[0]
    nsteps = len(attention_maps_all[rand_name])
    hw_ori = width * height

    # html_path = save_dir + '.html'
    text_input = save_dir.split('/')[-1]
    # f = open(html_path, 'wt')

    all_html_paths = []

    for step_num in range(0, nsteps, 5):

        # if cond_id == 'cond':
        #     attention_maps_cur = attention_maps_cond[step_num]
        # else:
        #     attention_maps_cur = attention_maps_uncond[step_num]

        attention_maps = dict()

        for layer in attention_maps_all.keys():

            attention_ind = attention_maps_all[layer][step_num].cpu()

            # Attention maps are of shape [batch_size, nkeys, 77]
            # since they are averaged out while collecting from hooks to save memory.
            # Now split the heads from batch dimension
            bs, hw, nclip = attention_ind.shape
            down_ratio = np.sqrt(hw_ori // hw)
            width_cur = int(width // down_ratio)
            height_cur = int(height // down_ratio)
            attention_ind = attention_ind.reshape(
                bs, height_cur, width_cur, nclip)

            attention_maps[layer] = attention_ind

        # Obtain heatmaps corresponding to random heads and individual heads

        html_names = save_attention_heatmaps(
            attention_maps, tokens_vis, save_dir=save_dir, prefix='step_{}/attention_maps_cond'.format(
                step_num)
        )

        # Write the logic for recursively creating pages
        for html_name_cur in html_names:
            all_html_paths.append(os.path.join(text_input, html_name_cur))

    save_dir_root = '/'.join(save_dir.split('/')[0:-1])
    for html_pth in all_html_paths:
        create_recursive_html_link(html_pth, save_dir_root)


def plot_attention_maps(atten_map_list, obj_tokens, save_dir, seed, tokens_vis=None):
    for i, attn_map in enumerate(atten_map_list):
        n_obj = len(attn_map)
        plt.figure()
        plt.clf()

        fig, axs = plt.subplots(
            ncols=n_obj+1, gridspec_kw=dict(width_ratios=[1 for _ in range(n_obj)]+[0.1]))

        fig.set_figheight(3)
        fig.set_figwidth(3*n_obj+0.1)

        cmap = plt.get_cmap('YlOrRd')

        vmax = 0
        vmin = 1
        for tid in range(n_obj):
            attention_map_cur = attn_map[tid]
            vmax = max(vmax, float(attention_map_cur.max()))
            vmin = min(vmin, float(attention_map_cur.min()))

        for tid in range(n_obj):
            sns.heatmap(
                attn_map[tid][0], annot=False, cbar=False, ax=axs[tid],
                cmap=cmap, vmin=vmin, vmax=vmax
            )
            axs[tid].set_axis_off()

            if tokens_vis is not None:
                if tid == n_obj-1:
                    axs_xlabel = 'other tokens'
                else:
                    axs_xlabel = ''
                    for token_id in obj_tokens[tid]:
                        axs_xlabel += ' ' + tokens_vis[token_id.item() -
                                                       1][:-len('</w>')]
                axs[tid].set_title(axs_xlabel)

        norm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
        fig.colorbar(sm, cax=axs[-1])

        fig.tight_layout()

        canvas = fig.canvas
        canvas.draw()
        width, height = canvas.get_width_height()
        img = np.frombuffer(canvas.tostring_rgb(),
                            dtype='uint8').reshape((height, width, 3))
        plt.savefig(os.path.join(
            save_dir, 'average_seed%d_attn%d.jpg' % (seed, i)), dpi=100)
        plt.close('all')
    return img


def get_average_attention_maps(attention_maps, save_dir, width, height, obj_tokens, seed=0, tokens_vis=None,
                               preprocess=False):
    r"""Function to visualize attention maps.
    Args:
        save_dir (str): Path to save attention maps
        batch_size (int): Batch size
        sampler_order (int): Sampler order
    """

    # Split attention maps over steps
    attention_maps_cond, _ = split_attention_maps_over_steps(
        attention_maps
    )

    nsteps = len(attention_maps_cond)
    hw_ori = width * height

    attention_maps = []
    for obj_token in obj_tokens:
        attention_maps.append([])

    for step_num in range(nsteps):
        attention_maps_cur = attention_maps_cond[step_num]

        for layer in attention_maps_cur.keys():
            if step_num < 10 or layer not in CrossAttentionLayers:
                continue

            attention_ind = attention_maps_cur[layer].cpu()

            # Attention maps are of shape [batch_size, nkeys, 77]
            # since they are averaged out while collecting from hooks to save memory.
            # Now split the heads from batch dimension
            bs, hw, nclip = attention_ind.shape
            down_ratio = np.sqrt(hw_ori // hw)
            width_cur = int(width // down_ratio)
            height_cur = int(height // down_ratio)
            attention_ind = attention_ind.reshape(
                bs, height_cur, width_cur, nclip)
            for obj_id, obj_token in enumerate(obj_tokens):
                if obj_token[0] == -1:
                    attention_map_prev = torch.stack(
                        [attention_maps[i][-1] for i in range(obj_id)]).sum(0)
                    attention_maps[obj_id].append(
                        attention_map_prev.max()-attention_map_prev)
                else:
                    obj_attention_map = attention_ind[:, :, :, obj_token].max(-1, True)[
                        0].permute([3, 0, 1, 2])
                    # obj_attention_map = attention_ind[:, :, :, obj_token].mean(-1, True).permute([3, 0, 1, 2])
                    obj_attention_map = torchvision.transforms.functional.resize(obj_attention_map, (height, width),
                                                                                 interpolation=torchvision.transforms.InterpolationMode.BICUBIC, antialias=True)
                    attention_maps[obj_id].append(obj_attention_map)

    attention_maps_averaged = []
    for obj_id, obj_token in enumerate(obj_tokens):
        if obj_id == len(obj_tokens) - 1:
            attention_maps_averaged.append(
                torch.cat(attention_maps[obj_id]).mean(0))
        else:
            attention_maps_averaged.append(
                torch.cat(attention_maps[obj_id]).mean(0))

    attention_maps_averaged_normalized = []
    attention_maps_averaged_sum = torch.cat(attention_maps_averaged).sum(0)
    for obj_id, obj_token in enumerate(obj_tokens):
        attention_maps_averaged_normalized.append(
            attention_maps_averaged[obj_id]/attention_maps_averaged_sum)

    if obj_tokens[-1][0] != -1:
        attention_maps_averaged_normalized = (
            torch.cat(attention_maps_averaged)/0.001).softmax(0)
        attention_maps_averaged_normalized = [
            attention_maps_averaged_normalized[i:i+1] for i in range(attention_maps_averaged_normalized.shape[0])]

    if preprocess:
        selem = square(5)
        selem = square(3)
        selem = square(1)
        attention_maps_averaged_eroded = [erosion(skimage.img_as_float(
            map[0].numpy()*255), selem) for map in attention_maps_averaged_normalized[:2]]
        attention_maps_averaged_eroded = [(torch.from_numpy(map).unsqueeze(
            0)/255. > 0.8).float() for map in attention_maps_averaged_eroded]
        attention_maps_averaged_eroded.append(
            1 - torch.cat(attention_maps_averaged_eroded).sum(0, True))
        plot_attention_maps([attention_maps_averaged, attention_maps_averaged_normalized,
                            attention_maps_averaged_eroded], obj_tokens, save_dir, seed, tokens_vis)
        attention_maps_averaged_eroded = [attn_mask.unsqueeze(1).repeat(
            [1, 4, 1, 1]).cuda() for attn_mask in attention_maps_averaged_eroded]
        return attention_maps_averaged_eroded
    else:
        plot_attention_maps([attention_maps_averaged, attention_maps_averaged_normalized],
                            obj_tokens, save_dir, seed, tokens_vis)
        attention_maps_averaged_normalized = [attn_mask.unsqueeze(1).repeat(
            [1, 4, 1, 1]).cuda() for attn_mask in attention_maps_averaged_normalized]
        return attention_maps_averaged_normalized


def get_average_attention_maps_threshold(attention_maps, save_dir, width, height, obj_tokens, seed=0, threshold=0.02):
    r"""Function to visualize attention maps.
    Args:
        save_dir (str): Path to save attention maps
        batch_size (int): Batch size
        sampler_order (int): Sampler order
    """

    _EPS = 1e-8
    # Split attention maps over steps
    attention_maps_cond, _ = split_attention_maps_over_steps(
        attention_maps
    )

    nsteps = len(attention_maps_cond)
    hw_ori = width * height

    attention_maps = []
    for obj_token in obj_tokens:
        attention_maps.append([])

    # for each side prompt, get attention maps for all steps and all layers
    for step_num in range(nsteps):
        attention_maps_cur = attention_maps_cond[step_num]
        for layer in attention_maps_cur.keys():
            attention_ind = attention_maps_cur[layer].cpu()
            bs, hw, nclip = attention_ind.shape
            down_ratio = np.sqrt(hw_ori // hw)
            width_cur = int(width // down_ratio)
            height_cur = int(height // down_ratio)
            attention_ind = attention_ind.reshape(
                bs, height_cur, width_cur, nclip)
            for obj_id, obj_token in enumerate(obj_tokens):
                if attention_ind.shape[1] > width//2:
                    continue
                if obj_token[0] != -1:
                    obj_attention_map = attention_ind[:, :, :,
                                                      obj_token].mean(-1, True).permute([3, 0, 1, 2])
                    obj_attention_map = torchvision.transforms.functional.resize(obj_attention_map, (height, width),
                                                                                 interpolation=torchvision.transforms.InterpolationMode.BICUBIC, antialias=True)
                    attention_maps[obj_id].append(obj_attention_map)

    # average of all steps and layers, thresholding
    attention_maps_thres = []
    attention_maps_averaged = []
    for obj_id, obj_token in enumerate(obj_tokens):
        if obj_token[0] != -1:
            average_map = torch.cat(attention_maps[obj_id]).mean(0)
            attention_maps_averaged.append(average_map)
            attention_maps_thres.append((average_map > threshold).float())

    # get the remaining region except for the original prompt
    attention_maps_averaged_normalized = []
    attention_maps_averaged_sum = torch.cat(attention_maps_thres).sum(0) + _EPS
    for obj_id, obj_token in enumerate(obj_tokens):
        if obj_token[0] != -1:
            attention_maps_averaged_normalized.append(
                attention_maps_thres[obj_id]/attention_maps_averaged_sum)
        else:
            attention_map_prev = torch.stack(
                attention_maps_averaged_normalized).sum(0)
            attention_maps_averaged_normalized.append(1.-attention_map_prev)

    plot_attention_maps(
        [attention_maps_averaged, attention_maps_averaged_normalized], save_dir, seed)

    attention_maps_averaged_normalized = [attn_mask.unsqueeze(1).repeat(
        [1, 4, 1, 1]).cuda() for attn_mask in attention_maps_averaged_normalized]
    # attention_maps_averaged_normalized = attention_maps_averaged_normalized.unsqueeze(1).repeat([1, 4, 1, 1]).cuda()
    return attention_maps_averaged_normalized


def get_token_maps(selfattn_maps, crossattn_maps, n_maps, save_dir, width, height, obj_tokens, kmeans_seed=0, tokens_vis=None,
                   preprocess=False, segment_threshold=0.3, num_segments=5, return_vis=False, save_attn=False):
    r"""Function to visualize attention maps.
    Args:
        save_dir (str): Path to save attention maps
        batch_size (int): Batch size
        sampler_order (int): Sampler order
    """

    resolution = 32
    # attn_maps_1024 = [attn_map for attn_map in selfattn_maps.values(
    # ) if attn_map.shape[1] == resolution**2]
    # attn_maps_1024 = torch.cat(attn_maps_1024).mean(0).cpu().numpy()
    attn_maps_1024 = {8: [], 16: [], 32: [], 64: []}
    for attn_map in selfattn_maps.values():
        resolution_map = np.sqrt(attn_map.shape[1]).astype(int)
        if resolution_map != resolution:
            continue
        # attn_map = torch.nn.functional.interpolate(rearrange(attn_map, '1 c (h w) -> 1 c h w', h=resolution_map), (resolution, resolution),
        #                                            mode='bicubic', antialias=True)
        # attn_map = rearrange(attn_map, '1 (h w) a b -> 1 (a b) h w', h=resolution_map)
        attn_map = attn_map.reshape(
            1, resolution_map, resolution_map, resolution_map**2).permute([3, 0, 1, 2]).float()
        attn_map = torch.nn.functional.interpolate(attn_map, (resolution, resolution),
                                                   mode='bicubic', antialias=True)
        attn_maps_1024[resolution_map].append(attn_map.permute([1, 2, 3, 0]).reshape(
            1, resolution**2, resolution_map**2))
    attn_maps_1024 = torch.cat([torch.cat(v).mean(0).cpu()
                                for v in attn_maps_1024.values() if len(v) > 0], -1).numpy()
    if save_attn:
        print('saving self-attention maps...', attn_maps_1024.shape)
        torch.save(torch.from_numpy(attn_maps_1024),
                   'results/maps/selfattn_maps.pth')
    seed_everything(kmeans_seed)
    # import ipdb;ipdb.set_trace()
    # kmeans = KMeans(n_clusters=num_segments,
    #                 n_init=10).fit(attn_maps_1024)
    # clusters = kmeans.labels_
    # clusters = clusters.reshape(resolution, resolution)
    # mesh = np.array(np.meshgrid(range(resolution), range(resolution), indexing='ij'), dtype=np.float32)/resolution
    # dists = mesh.reshape(2, -1).T
    # delta = 0.01
    # spatial_sim = rbf_kernel(dists, dists)*delta
    sc = SpectralClustering(num_segments, affinity='precomputed', n_init=100,
                            assign_labels='kmeans')
    clusters = sc.fit_predict(attn_maps_1024)
    clusters = clusters.reshape(resolution, resolution)
    fig = plt.figure()
    plt.imshow(clusters)
    plt.axis('off')
    plt.savefig(os.path.join(save_dir, 'segmentation_k%d_seed%d.jpg' % (num_segments, kmeans_seed)),
                bbox_inches='tight', pad_inches=0)
    if return_vis:
        canvas = fig.canvas
        canvas.draw()
        cav_width, cav_height = canvas.get_width_height()
        segments_vis = np.frombuffer(canvas.tostring_rgb(),
                                     dtype='uint8').reshape((cav_height, cav_width, 3))

    plt.close()

    # label the segmentation mask using cross-attention maps
    cross_attn_maps_1024 = []
    for attn_map in crossattn_maps.values():
        resolution_map = np.sqrt(attn_map.shape[1]).astype(int)
        # if resolution_map != 16:
        # continue
        attn_map = attn_map.reshape(
            1, resolution_map, resolution_map, -1).permute([0, 3, 1, 2]).float()
        attn_map = torch.nn.functional.interpolate(attn_map, (resolution, resolution),
                                                   mode='bicubic', antialias=True)
        cross_attn_maps_1024.append(attn_map.permute([0, 2, 3, 1]))

    cross_attn_maps_1024 = torch.cat(
        cross_attn_maps_1024).mean(0).cpu().numpy()
    normalized_span_maps = []
    for token_ids in obj_tokens:
        token_ids = [token_id for token_id in token_ids if token_id < 77]
        span_token_maps = cross_attn_maps_1024[:, :, token_ids.numpy()]
        normalized_span_map = np.zeros_like(span_token_maps)
        for i in range(span_token_maps.shape[-1]):
            curr_noun_map = span_token_maps[:, :, i]
            normalized_span_map[:, :, i] = (
                # curr_noun_map - np.abs(curr_noun_map.min())) / curr_noun_map.max()
                curr_noun_map - np.abs(curr_noun_map.min())) / (curr_noun_map.max()-curr_noun_map.min())
        normalized_span_maps.append(normalized_span_map)
    foreground_token_maps = [np.zeros([clusters.shape[0], clusters.shape[1]]).squeeze(
    ) for normalized_span_map in normalized_span_maps]
    background_map = np.zeros([clusters.shape[0], clusters.shape[1]]).squeeze()
    for c in range(num_segments):
        cluster_mask = np.zeros_like(clusters)
        cluster_mask[clusters == c] = 1.
        is_foreground = False
        for normalized_span_map, foreground_nouns_map, token_ids in zip(normalized_span_maps, foreground_token_maps, obj_tokens):
            score_maps = [cluster_mask * normalized_span_map[:, :, i]
                          for i in range(len(token_ids))]
            scores = [score_map.sum() / cluster_mask.sum()
                      for score_map in score_maps]
            if max(scores) > segment_threshold:
                foreground_nouns_map += cluster_mask
                is_foreground = True
        if not is_foreground:
            background_map += cluster_mask
    foreground_token_maps.append(background_map)

    # resize the token maps and visualization
    resized_token_maps = torch.cat([torch.nn.functional.interpolate(torch.from_numpy(token_map).unsqueeze(0).unsqueeze(
        0), (height, width), mode='bicubic', antialias=True)[0] for token_map in foreground_token_maps]).clamp(0, 1)

    resized_token_maps = resized_token_maps / \
        (resized_token_maps.sum(0, True)+1e-8)
    resized_token_maps = [token_map.unsqueeze(
        0) for token_map in resized_token_maps]
    foreground_token_maps = [token_map[None, :, :]
                             for token_map in foreground_token_maps]
    if preprocess:
        selem = square(5)
        eroded_token_maps = torch.stack([torch.from_numpy(erosion(skimage.img_as_float(
            map[0].numpy()*255), selem))/255. for map in resized_token_maps[:-1]]).clamp(0, 1)
        # import ipdb; ipdb.set_trace()
        eroded_background_maps = (1-eroded_token_maps.sum(0, True)).clamp(0, 1)
        eroded_token_maps = torch.cat([eroded_token_maps, eroded_background_maps])
        eroded_token_maps = eroded_token_maps / (eroded_token_maps.sum(0, True)+1e-8)
        resized_token_maps = [token_map.unsqueeze(
            0) for token_map in eroded_token_maps]

    token_maps_vis = plot_attention_maps([foreground_token_maps, resized_token_maps], obj_tokens,
                                         save_dir, kmeans_seed, tokens_vis)
    resized_token_maps = [token_map.unsqueeze(1).repeat(
        [1, 4, 1, 1]).to(attn_map.dtype).cuda() for token_map in resized_token_maps]
    if return_vis:
        return resized_token_maps, segments_vis, token_maps_vis
    else:
        return resized_token_maps