muse-maskgit-pytorch

Why we included this project

Google's Muse generates images from text with masked generative transformers, and this repo is a PyTorch re-implementation of that approach, organized as components you train rather than a ready-made pipeline. There is a VQGan VAE that tokenizes images into a codebook, a MaskGit transformer that predicts masked image tokens from a text prompt, and a Muse wrapper that chains a low-resolution base model with a super-resolution stage to reach 256 or 512 pixel outputs. Codebook size, transformer depth, and classifier-free guidance scale are all exposed, so you can train on small images first and enlarge later. That makes it a practical base for studying how token-based text-to-image generation works, or for training custom generators on your own datasets instead of consuming pretrained checkpoints. The end-to-end training and generation code is all present, which saves the trouble of porting the paper yourself.