magvit2 pytorch

MagViT2 - Pytorch

A implementação de MagViT2 do Language Model Beats Diffusion - Tokenizer é a chave para a geração visual em Pytorch. Atualmente, isso contém SOTA para geração/compreensão de vídeo.

O Lookup Free Quantizer proposto no artigo pode ser encontrado em um repositório separado. Provavelmente deveria ser explorado para todas as outras modalidades, começando com áudio

Por favor, junte-se se estiver interessado em replicar o tokenizer proposto neste artigo abertamente

Atualização: Tencent usou o código neste repositório e abriu o código-fonte de um modelo funcional

Apreciação

• EstabilidadeAI e ? Huggingface pelo generoso patrocínio, assim como aos meus outros patrocinadores, por me proporcionarem independência para inteligência artificial de código aberto.

• Louis Serrano por compartilhar algumas execuções iniciais, validando que a arquitetura geral converge com a quantização escalar finita.

• Você? Se você é um engenheiro/cientista talentoso, sinta-se à vontade para contribuir com a ciência de ponta em código aberto!

Instalar

$ pip install magvit2-pytorch

Uso

 from magvit2_pytorch import (
    VideoTokenizer ,
    VideoTokenizerTrainer
)

tokenizer = VideoTokenizer (
    image_size = 128 ,
    init_dim = 64 ,
    max_dim = 512 ,
    codebook_size = 1024 ,
    layers = (
        'residual' ,
        'compress_space' ,
        ( 'consecutive_residual' , 2 ),
        'compress_space' ,
        ( 'consecutive_residual' , 2 ),
        'linear_attend_space' ,
        'compress_space' ,
        ( 'consecutive_residual' , 2 ),
        'attend_space' ,
        'compress_time' ,
        ( 'consecutive_residual' , 2 ),
        'compress_time' ,
        ( 'consecutive_residual' , 2 ),
        'attend_time' ,
    )
)

trainer = VideoTokenizerTrainer (
    tokenizer ,
    dataset_folder = '/path/to/a/lot/of/media' ,     # folder of either videos or images, depending on setting below
    dataset_type = 'videos' ,                        # 'videos' or 'images', prior papers have shown pretraining on images to be effective for video synthesis
    batch_size = 4 ,
    grad_accum_every = 8 ,
    learning_rate = 2e-5 ,
    num_train_steps = 1_000_000
)

trainer . train ()

# after a lot of training ...
# can use the EMA of the tokenizer

ema_tokenizer = trainer . ema_tokenizer

# mock video

video = torch . randn ( 1 , 3 , 17 , 128 , 128 )

# tokenizing video to discrete codes

codes = ema_tokenizer . tokenize ( video ) # (1, 9, 16, 16) <- in this example, time downsampled by 4x and space downsampled by 8x. flatten token ids for (non)-autoregressive training

# sanity check

decoded_video = ema_tokenizer . decode_from_code_indices ( codes )

assert torch . allclose (
    decoded_video ,
    ema_tokenizer ( video , return_recon = True )
)

Para rastrear seus experimentos em Pesos e preconceitos, defina use_wandb_tracking = True em VideoTokenizerTrainer e, em seguida, use o gerenciador de contexto .trackers

 trainer = VideoTokenizerTrainer (
    use_wandb_tracking = True ,
    ...
)

with trainer . trackers ( project_name = 'magvit2' , run_name = 'baseline' ):
    trainer . train ()

Pendência

• Tokenizador Magvit2

  • adicionar perda adversária
  • implementar o blurpool para antialiasing no discriminador
  • O LFQ deve ser capaz de passar a quebra de perdas (comprometimento e entropia) e encaminhar para o retorno do tokenizer
  • adicionar condicionamento para decodificador de codificador com conversão modulável residual 3d
  • decode_from_codebook_indices deve ser capaz de aceitar IDs achatados e remodelar para corrigir as dimensões do mapa de recursos e decodificar de volta para o vídeo
  • adicionar treinador e gerenciar treinamento discriminador
  • adicionar rmsnorm adaptativo e camadas de transformador condicionais
  • generalizar completamente para múltiplos discriminadores em diferentes escalas de tempo (inspirando-se em discriminadores de múltiplas resoluções do fluxo sonoro)
    • perdas completas do discriminador multiescala
    • gerenciar automaticamente otimizadores discriminadores multiescala
    • funções auxiliares para criar discriminadores temporais de multi-resolução (escolher quadros consecutivos aleatórios)
  • adicionar atenção
    • usar embutimentos rotativos axiais para espacial
  • adicione uma perda autoregressiva opcional em alguma penúltima camada do decodificador - verifique a literatura para ver se alguém mais fez essa unificação do decodificador do transformador + tokenizador em uma arquitetura

• Improvise um transformador de vídeo RQ, já que o LFQ residual realmente faz sentido agora

• MáscaraGit

Citações

 @misc { yu2023language ,
    title   = { Language Model Beats Diffusion -- Tokenizer is Key to Visual Generation } , 
    author  = { Lijun Yu and José Lezama and Nitesh B. Gundavarapu and Luca Versari and Kihyuk Sohn and David Minnen and Yong Cheng and Agrim Gupta and Xiuye Gu and Alexander G. Hauptmann and Boqing Gong and Ming-Hsuan Yang and Irfan Essa and David A. Ross and Lu Jiang } ,
    year    = { 2023 } ,
    eprint  = { 2310.05737 } ,
    archivePrefix = { arXiv } ,
    primaryClass = { cs.CV }
}

 @inproceedings { dao2022flashattention ,
    title   = { Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness } ,
    author  = { Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{'e}, Christopher } ,
    booktitle = { Advances in Neural Information Processing Systems } ,
    year    = { 2022 }
}

 @article { Zhang2021TokenST ,
    title   = { Token Shift Transformer for Video Classification } ,
    author  = { Hao Zhang and Y. Hao and Chong-Wah Ngo } ,
    journal = { Proceedings of the 29th ACM International Conference on Multimedia } ,
    year    = { 2021 }
}

 @inproceedings { Arora2023ZoologyMA ,
    title   = { Zoology: Measuring and Improving Recall in Efficient Language Models } ,
    author  = { Simran Arora and Sabri Eyuboglu and Aman Timalsina and Isys Johnson and Michael Poli and James Zou and Atri Rudra and Christopher R'e } ,
    year    = { 2023 } ,
    url     = { https://api.semanticscholar.org/CorpusID:266149332 }
}