Function aggregate_attention_all_step

sample.py:230–255

(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int, len_tokens: int)

def aggregate_attention_all_step(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int, len_tokens: int):
    attention_all_step_maps = attention_store.all_step_attention

    out_all_step = {}
    for step, attention_maps in attention_all_step_maps.items():
        out = []
        for _, value in attention_maps.items():
            for item in value:
                res = int(math.sqrt(item.shape[1]))
                cross_maps = item.reshape(len(prompts), -1, res , res, item.shape[-1])[select]
                cross_maps_512 = torch.zeros((cross_maps.shape[0], 256, 256, 77))
                for token in range(len_tokens):
                    token_cross_map = cross_maps[:, :, :, token].reshape(1, cross_maps.shape[0], res, res)
                    token_cross_map = F.interpolate(token_cross_map, size=(256, 256), mode='bicubic').reshape(cross_maps.shape[0], 256, 256)
                    cross_maps_512[:, :, :, token] = token_cross_map[:, :, :]
                out.append(cross_maps_512)
        print(step)
        out = torch.cat(out, dim=0)
        out = out.sum(0) / out.shape[0]

        absolute = True
        if not absolute:
            out = (out - out.min()) / (out.max() - out.min() + 1e-8)
        out_all_step[step] = out.cpu()

    return out_all_step

def show_cross_attention(prompts, tokenizer, save_path, base_count, attention_store: AttentionStore, res: int, from_where: List[str], select: int = 0):
    tokens = tokenizer.encode(prompts[select])