Function aggregate_attention

sample.py:189–200

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

    # display(pil_img)

def aggregate_attention(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int):
    out = []
    attention_maps = attention_store.get_average_attention()
    num_pixels = res ** 2
    for location in from_where:
        for item in attention_maps[f"{location}_{'cross' if is_cross else 'self'}"]:
            if item.shape[1] == num_pixels:
                cross_maps = item.reshape(len(prompts), -1, res, res, item.shape[-1])[select]
                out.append(cross_maps)
    out = torch.cat(out, dim=0)
    out = out.sum(0) / out.shape[0]
    return out.cpu()

def aggregate_attention_all(prompts, attention_store: AttentionStore, res: int, from_where: List[str], is_cross: bool, select: int, len_tokens: int):
    attention_maps = attention_store.get_average_attention()