| """ |
| 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() |
|
|
