Function main

tokenize_examples/tokenize_video_muse.py:51–93

(argv)

def main(argv):
    assert FLAGS.input_dir != ''
    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()

    videos = []
    for root, _, files in os.walk(FLAGS.input_dir):
        for file in files:
            if is_video(file):
                videos.append(os.path.join(root, file))

    with open(FLAGS.output_file, 'w') as fout:
        with torch.no_grad():
            for epoch in trange(FLAGS.n_epochs, ncols=0):
                for stride in tqdm(FLAGS.strides.split(','), ncols=0):
                    stride = int(stride)
                    dataset = VideoDataset(videos, n_frames=FLAGS.n_frames, stride=stride)
                    dataloader = torch.utils.data.DataLoader(
                        dataset,
                        batch_size=FLAGS.batch_size,
                        shuffle=False,
                        num_workers=FLAGS.n_workers,
                        prefetch_factor=4,
                        drop_last=True,
                    )
                    for batch in tqdm(dataloader, ncols=0):
                        batch_size = batch.shape[0]
                        batch = einops.rearrange(
                            batch.numpy(), 'b t h w c -> (b t) c h w'
                        )
                        batch = torch.tensor(batch).to(device)
                        _, tokens = net.encode(batch)
                        tokens = einops.rearrange(
                            tokens.cpu().numpy().astype(np.int32), '(b t) d -> b (t d)', b=batch_size
                        )
                        for i in range(batch_size):
                            data = {'tokens': b64encode(tokens[i].tobytes()).decode('utf-8'),}
                            fout.write(json.dumps(data) + '\n')