| """ | |
| This file is used to extract feature of the coco val set (to test zero-shot FID). | |
| """ | |
| import os | |
| import sys | |
| sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))) | |
| import torch | |
| import os | |
| import numpy as np | |
| from datasets import MSCOCODatabase | |
| import argparse | |
| from tqdm import tqdm | |
| import libs.autoencoder | |
| from libs.clip import FrozenCLIPEmbedder | |
| from libs.t5 import T5Embedder | |
| def main(resolution=256): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--split', default='val') | |
| args = parser.parse_args() | |
| print(args) | |
| if args.split == "val": | |
| datas = MSCOCODatabase(root='/data/qihao/dataset/coco2014/val2014', | |
| annFile='/data/qihao/dataset/coco2014/annotations/captions_val2014.json', | |
| size=resolution) | |
| save_dir = f'val' | |
| else: | |
| raise NotImplementedError | |
| device = "cuda" | |
| os.makedirs(save_dir, exist_ok=True) | |
| autoencoder = libs.autoencoder.get_model('../assets/stable-diffusion/autoencoder_kl.pth') | |
| autoencoder.to(device) | |
| llm = 'clip' | |
| if llm=='clip': | |
| clip = FrozenCLIPEmbedder() | |
| clip.eval() | |
| clip.to(device) | |
| elif llm=='t5': | |
| t5 = T5Embedder(device=device) | |
| else: | |
| raise NotImplementedError | |
| with torch.no_grad(): | |
| for idx, data in tqdm(enumerate(datas)): | |
| x, captions = data | |
| if len(x.shape) == 3: | |
| x = x[None, ...] | |
| x = torch.tensor(x, device=device) | |
| moments = autoencoder(x, fn='encode_moments').squeeze(0) | |
| moments = moments.detach().cpu().numpy() | |
| np.save(os.path.join(save_dir, f'{idx}.npy'), moments) | |
| if llm=='clip': | |
| latent, latent_and_others = clip.encode(captions) | |
| token_embedding = latent_and_others['token_embedding'] | |
| token_mask = latent_and_others['token_mask'] | |
| token = latent_and_others['tokens'] | |
| elif llm=='t5': | |
| latent, latent_and_others = t5.get_text_embeddings(captions) | |
| token_embedding = latent_and_others['token_embedding'].to(torch.float32) * 10.0 | |
| token_mask = latent_and_others['token_mask'] | |
| token = latent_and_others['tokens'] | |
| for i in range(len(captions)): | |
| data = {'promt': captions[i], | |
| 'token_embedding': token_embedding[i].detach().cpu().numpy(), | |
| 'token_mask': token_mask[i].detach().cpu().numpy(), | |
| 'token': token[i].detach().cpu().numpy()} | |
| np.save(os.path.join(save_dir, f'{idx}_{i}.npy'), data) | |
| if __name__ == '__main__': | |
| main() | |