Function show_cross_attention

sample.py:257–295

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

    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])
    decoder = tokenizer.decode
    
    attention_maps = aggregate_attention_all(prompts, attention_store, res, from_where, True, select, len(tokens))
    images = []
    for i in range(len(tokens)):
        image = attention_maps[:, :, i]
        image = 255 * image / image.max()
        image[image < 0] = 0 # process the overflow
        image = image.unsqueeze(-1).expand(*image.shape, 3)
        image = image.numpy().astype(np.uint8)
        image = np.array(Image.fromarray(image).resize((256, 256)))
        # tmp_image = Image.fromarray(image) 
        # tmp_image.save(os.path.join(save_path, '{}_{}.png'.format(base_count, decoder(int(tokens[i])))))
        text_under_image_string = decoder(int(tokens[i]))
        if text_under_image_string == '<new1>':
            text_under_image_string = 'V1*'
        elif text_under_image_string == '<new2>':
            text_under_image_string = 'V2*'
        image = text_under_image(image, text_under_image_string)
        # image = text_under_image(image, decoder(int(tokens[i])))

        images.append(image)
    view_images(np.stack(images, axis=0), save_path, base_count)

    if attention_store.all_step_attention_store:
        attention_maps_all_step = aggregate_attention_all_step(prompts, attention_store, res, from_where, True, select, len(tokens))
        for step, attention_maps in attention_maps_all_step.items():
            images = []
            for i in range(len(tokens)):
                image = attention_maps[:, :, i]
                image = 255 * image / image.max()
                image = image.unsqueeze(-1).expand(*image.shape, 3)
                image = image.numpy().astype(np.uint8)
                image = np.array(Image.fromarray(image).resize((256, 256)))
                image = text_under_image(image, decoder(int(tokens[i])))
                images.append(image)
            view_images(np.stack(images, axis=0), save_path, base_count, step=step)


def main():