DOCSEARCHbyAlgolia
Submit Docs

muse-maskgit-pytorch

repository·main·Indexed 21 days ago

https://github.com/lucidrains/muse-maskgit-pytorch

A PyTorch implementation of the Muse model for text-to-image generation using Masked Generative Transformers. It utilizes a VQ-GAN VAE and a transformer-based MaskGit approach to generate high-quality images, supporting both base text-to-image generation and super-resolution pipelines via the Muse class.

Tokens
1.8K
Snippets
5
Records
5
Agent score
26%

What's inside muse-maskgit-pytorch

  2. Use MaskGit for Text-to-Image Generation

    main

    To use MaskGit, you need a trained VQGanVAE and a MaskGitTransformer.

    Key Requirements:

    • MaskGitTransformer.num_tokens must match VQGanVAE.codebook_size.
    • MaskGitTransformer.seq_len must be equivalent to fmap_size ** 2 from the VAE.

    Training: Pass images and texts to the MaskGit instance with return_loss=True to compute the loss.

    Generation: Use the .generate() method with texts and an optional cond_scale for classifier-free guidance.

    import torch
from muse_maskgit_pytorch import VQGanVAE, MaskGit, MaskGitTransformer

# 1. Instantiate VAE
vae = VQGanVAE(
    dim = 256,
    codebook_size = 65536
).cuda()
vae.load('/path/to/vae.pt')

# 2. Create Transformer
transformer = MaskGitTransformer(
    num_tokens = 65536,       # must be same as codebook size above
    seq_len = 256,            # must be equivalent to fmap_size ** 2 in vae
    dim = 512,
    depth = 8,
    dim_head = 64,
    heads = 8,
    ff_mult = 4,
    t5_name = 't5-small',
)

# 3. Instantiate MaskGit
base_maskgit = MaskGit(
    vae = vae,
    transformer = transformer,
    image_size = 256,
    cond_drop_prob = 0.25,
).cuda()

# Training step
texts = ['a child screaming at finding a worm within a half-eaten apple']
images = torch.randn(4, 3, 256, 256).cuda()
loss = base_maskgit(images, texts = texts)
loss.backward()

# Generation
images = base_maskgit.generate(
    texts = ['a whale breaching from afar'],
    cond_scale = 3.
)

  3. Use MaskGit for Super-Resolution

    main

    To perform super-resolution, configure MaskGit with cond_image_size (the size of the low-resolution conditioning image) and a larger image_size (the target high-resolution size).

    Training: Pass the high-resolution images and texts to the model.

    Generation: When calling .generate(), you must provide cond_images (the low-resolution versions of the target images).

    import torch
import torch.nn.functional as F
from muse_maskgit_pytorch import VQGanVAE, MaskGit, MaskGitTransformer

vae = VQGanVAE(dim = 256, codebook_size = 65536).cuda()
vae.load('./path/to/vae.pt')

transformer = MaskGitTransformer(
    num_tokens = 65536,
    seq_len = 1024,
    dim = 512,
    depth = 2,
    dim_head = 64,
    heads = 8,
    ff_mult = 4,
    t5_name = 't5-small',
)

superres_maskgit = MaskGit(
    vae = vae,
    transformer = transformer,
    cond_drop_prob = 0.25,
    image_size = 512,      # target high-res size
    cond_image_size = 256, # conditioning low-res size
).cuda()

# Training
images = torch.randn(4, 3, 512, 512).cuda()
texts = ['a child screaming at finding a worm within a half-eaten apple']
loss = superres_maskgit(images, texts = texts)
loss.backward()

# Generation
images = superres_maskgit.generate(
    texts = ['a whale breaching from afar'],
    cond_images = F.interpolate(images, 256), # must pass conditioning images
    cond_scale = 3.
)

  4. Train a VQGanVAE

    main

    Before training the MaskGit model, you must first train a VAE (specifically VQGanVAE). Use the VQGanVAETrainer to train on a folder of images. It is recommended to start with smaller image sizes and use curriculum learning for larger ones.

    import torch
from muse_maskgit_pytorch import VQGanVAE, VQGanVAETrainer

vae = VQGanVAE(
    dim = 256,
    codebook_size = 65536
)

# train on folder of images, as many images as possible

trainer = VQGanVAETrainer(
    vae = vae,
    image_size = 128,             # start small, then curriculum learn to larger ones
    folder = '/path/to/images',
    batch_size = 4,
    grad_accum_every = 8,
    num_train_steps = 50000
).cuda()

trainer.train()

  5. Use the Muse class for end-to-end generation

    main

    The Muse class combines a base MaskGit model and a super-resolution MaskGit model to provide a complete pipeline. Pass both trained models to Muse to generate high-resolution images from text prompts.

    from muse_maskgit_pytorch import Muse

base_maskgit.load('./path/to/base.pt')
superres_maskgit.load('./path/to/superres.pt')

# Pass in the trained base and superres models
muse = Muse(
    base = base_maskgit,
    superres = superres_maskgit
)

images = muse([
    'a whale breaching from afar',
    'young girl blowing out candles on her birthday cake'
])

# images is a List[PIL.Image.Image]