Function main

tokenize_examples/tokenize_multi_datasets_muse.py:111–176

(argv)

def main(argv):
    assert FLAGS.input_dir != ''
    assert FLAGS.input_regex != ''
    assert FLAGS.output_file != ''

    # Load the pre-trained vq model from the hub
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    net = VQGANModel.from_pretrained('vqlm/muse/ckpts/laion').to(device)
    net.eval()

    regex = FLAGS.input_regex.split('::')
    input_images = match_mulitple_path_v2(FLAGS.input_dir, regex)

    print(f'Found {len(input_images)} images')
    assert len(input_images) > 0, 'No images found'

    if FLAGS.max_examples > 0:
        input_images = input_images[:FLAGS.max_examples]

    random.shuffle(input_images)

    dataset = MultipleImageDataset(input_images)
    dataloader = torch.utils.data.DataLoader(
        dataset,
        batch_size=FLAGS.batch_size * FLAGS.n_shots,
        shuffle=False,
        num_workers=FLAGS.n_workers,
        drop_last=True
    )

    total_images = len(input_images) - len(input_images) % (FLAGS.batch_size * FLAGS.n_shots)

    with NamedTemporaryFile() as ntf:
        all_tokens = np.memmap(ntf, dtype='i4', mode='w+', shape=(total_images, 256 * len(input_images[0])))
        all_tokens[:] = 0

        index = 0
        for batch in tqdm(dataloader, ncols=0):
            k = 0
            for image in batch:
                batch_size = image.shape[0]
                image  = einops.rearrange(
                    image.numpy(), 'b h w c -> b c h w'
                )
                image  = torch.tensor(image).to(device)
                _, tokens = net.encode(image)
                tokens = einops.rearrange(
                    tokens.cpu().numpy().astype(np.int32), '(b t) d -> b (t d)', b=batch_size
                )
                all_tokens[index:index + image.shape[0], k:k + 256] = tokens
                k += 256
            index += batch[0].shape[0]

        with open(FLAGS.output_file, 'w') as fout:
            for _ in trange(FLAGS.n_epochs, ncols=0):
                indices = np.random.permutation(total_images).reshape(-1, FLAGS.n_shots)
                for i in trange(indices.shape[0], ncols=0):