Machinelearning

🌟

VidTok: Универсальный токенизатор видео от 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 чекпоинтов, расшифровка нейминга:

🟢

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