magvit2 pytorch

MagViT2-Pytorch

Implémentation de MagViT2 à partir du modèle de langage Beats Diffusion - Tokenizer est la clé de la génération visuelle dans Pytorch. Cela contient actuellement SOTA pour la génération/compréhension de vidéos.

Le Lookup Free Quantizer proposé dans l'article peut être trouvé dans un référentiel séparé. Il faudrait probablement l'explorer pour toutes les autres modalités, à commencer par l'audio.

Veuillez nous rejoindre si vous souhaitez reproduire à l'air libre le tokenizer proposé dans cet article.

Mise à jour : Tencent a utilisé le code de ce référentiel et a ouvert un modèle de travail

Appréciation

• StabilitéIA et ? Huggingface pour mon généreux parrainage, ainsi que mes autres sponsors, pour m'avoir offert l'indépendance nécessaire à l'intelligence artificielle open source.

• Louis Serrano pour avoir partagé quelques premières exécutions, validant que l'architecture globale converge avec une quantification scalaire finie.

• Toi? Si vous êtes un ingénieur de recherche/scientifique talentueux, n'hésitez pas à contribuer à la science open source de pointe !

Installer

$ pip install magvit2-pytorch

Usage

 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 )
)

Pour suivre vos expériences sur les poids et les biais, définissez use_wandb_tracking = True sur VideoTokenizerTrainer , puis utilisez le gestionnaire de contexte .trackers

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

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

Faire

• Tokeniseur Magvit2

  • ajouter une perte contradictoire
  • implémenter le Blupool pour l'anticrénelage dans le discriminateur
  • LFQ devrait pouvoir transmettre la répartition des pertes (engagement et entropie) et être transmis au retour du tokenizer
  • ajouter un conditionnement pour le décodeur d'encodeur avec une conv modulable résiduelle 3d
  • decode_from_codebook_indices devrait être capable d'accepter les identifiants aplatis et de se remodeler pour corriger les dimensions de la carte des caractéristiques et décoder en vidéo
  • ajouter un formateur et gérer la formation des discriminateurs
  • ajouter des couches de transformateur adaptatif rmsnorm et conditionnables
  • généraliser complètement à plusieurs discriminateurs à différentes échelles de temps (en s'inspirant des discriminateurs multi-résolution du flux sonore)
    • pertes complètes du discriminateur multi-échelles
    • gérer automatiquement les optimiseurs de discriminateur multi-échelles
    • fonctions d'assistance pour la création de discriminateurs temporels multi-résolution (sélection d'images consécutives aléatoires)
  • ajouter de l'attention
    • utiliser des intégrations rotatives axiales pour l'espace
  • ajoutez une perte autorégressive facultative au niveau d'une avant-dernière couche du décodeur - vérifiez la littérature pour voir si quelqu'un d'autre a fait cette unification du décodeur de transformateur + du tokenizer dans une seule architecture

• Improvisez un transformateur vidéo RQ, car le LFQ résiduel a désormais un sens

• MasqueGit

Citations

 @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 }
}