VidTok: Универсальный токенизатор видео от Microsoft

VidTok (https://github.com/microsoft/VidTok) – универсальный и открытый видео токенизатор, демонстрирующий высокую производительность как в непрерывной, так и в дискретной токенизации.

Токенизация видео, преобразующая исходные данные в компактные латентные токены - важнейший шаг для задач генерации и понимания видео. VidTok предлагает ряд улучшений, которые позволяют ему превзойти существующие методы: модельную архитектуру, методы квантования и стратегии обучения.

В архитектуре VidTok пространственное и временное сэмплирование обрабатываются раздельно, используя 2D свертки для пространственных модулей и оператор AlphaBlender для временных, при этом сохраняя 3D свертки для слияния информации.

Для дискретной токенизации используется конечное скалярное квантование (FSQ), которое оптимизирует неявный кодовый словарь, улучшая стабильность обучения. Эффективность обучения достигается двухэтапной стратегией: предварительное обучение на видео с низким разрешением, а затем дообучение декодера на видео с высоким разрешением.

VidTok обучался на датасете видеоданных с разным разрешением (400 000 видео 480p и 10 000 видео 1080p). Производительность измерялась с использованием метрик PSNR, SSIM, LPIPS и FVD, результаты показали превосходство VidTok по сравнению с другими токенизаторами как в дискретной, так и в непрерывной токенизации.

При сравнении с MAGVIT-v2, OmniTokenizer, CV-VAE, Open-Sora и Cosmos-Tokenizer, VidTok достиг лучших показателей, с меньшим размером модели.

В открытый доступ опубликованы 12 чекпоинтов (https://huggingface.co/microsoft/VidTok/tree/main/checkpoints), расшифровка нейминга:

vidtok - базовое название;

kl или fsq - тип регуляризации и квантования латентного пространства;

causal или noncausal - тип обработки временной информации (покадрово или все кадры сразу);

488 или 41616 - компрессионное соотношение (VCR), которое определяет степень сжатия видео по времени, высоте и ширине. Например, 4x8x8 и 4x16x16;

4chn, 8chn или 16chn - количество каналов в латентном пространстве для непрерывных токенизаторов. Чем больше каналов - тем качественней видео;

262144, 32768 или 4096 - размер codebook для дискретных токенизаторов с использованием FSQ. Чем больше - тем точнее представлятся информация.

Локальная установка и пример запуска как для непрерывной, так и для дискретной токенизации и как для каузальных, так и для некаузальных моделей:

# Clone repo  
 git clone https://github.com/microsoft/VidTok 
 
 cd VidTok 
 
 # Create conda env 
 
 conda env create -f environment.yaml 
 
 conda activate vidtok 
 
 # Inference  
 
 import torch 
 
 from scripts.inference_evaluate import load_model_from_config 
 
 cfg_path = "configs/vidtok_kl_causal_488_4chn.yaml" 
 
 ckpt_path = "checkpoints/vidtok_kl_causal_488_4chn.ckpt" 
 
 is_causal = True 
 
 device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")     
 
 # load pre-trained model 
 
 model = load_model_from_config(cfg_path, ckpt_path) 
 
 model.to(device).eval() 
 
 # random input 
 
 num_frames = 17 if is_causal else 16 
 
 x_input = (torch.rand(1, 3, num_frames, 256, 256) * 2 - 1).to(device)  # [B, C, T, H, W], range -1~1 
 
 # model forward 
 
 _, x_recon, _ = model(x_input) 
 
 assert x_input.shape == x_recon.shape

Лицензирование: MIT License.

Набор моделей (https://huggingface.co/microsoft/VidTok)

Arxiv (https://arxiv.org/pdf/2412.13061)

GitHub (https://github.com/microsoft/VidTok)

Источник: github.com

Комментарии: