DOCSEARCHbyAlgolia
Submit Docs

parti-pytorch

repository·main·Indexed 19 days ago

https://github.com/lucidrains/parti-pytorch

A PyTorch implementation of Google's Parti, a pure attention-based text-to-image neural network. The library includes the Parti model and components for training a Transformer-based VQ-GAN VAE (VitVQGanVAE) for visual tokenization, as well as utilities for pre-encoding text using T5 for efficient training.

Tokens
1.4K
Snippets
5
Records
5
Agent score
19%

What's inside parti-pytorch

  1. Train a Transformer VQ-GAN VAE

    main

    Before using Parti, you must train a VitVQGanVAE. This component acts as the visual tokenizer. You can use the VQGanVAETrainer to manage the training process, specifying the model, image directory, training steps, learning rate, and batch size.

    from parti_pytorch import VitVQGanVAE, VQGanVAETrainer

vit_vae = VitVQGanVAE(
    dim = 256,               # dimensions
    image_size = 256,        # target image size
    patch_size = 16,         # size of the patches in the image attending to each other
    num_layers = 3           # number of layers
).cuda()

trainer = VQGanVAETrainer(
    vit_vae,
    folder = '/path/to/your/images',
    num_train_steps = 100000,
    lr = 3e-4,
    batch_size = 4,
    grad_accum_every = 8,
    amp = True
)

trainer.train()

  2. Initialize and train the Parti model

    main

    To use Parti, instantiate the Parti class by passing in a trained VitVQGanVAE. You can then perform training by passing text prompts and images to the model instance with return_loss = True.

    import torch
from parti_pytorch import Parti, VitVQGanVAE

# 1. Instantiate VAE
vit_vae = VitVQGanVAE(
    dim = 256,
    image_size = 256,
    patch_size = 16,
    num_layers = 3
).cuda()

# Load your trained weights (preferably EMA)
vit_vae.load_state_dict(torch.load(f'/path/to/vae.pt'))

# 2. Instantiate Parti
parti = Parti(
    vae = vit_vae,
    dim = 512,
    depth = 8,
    dim_head = 64,
    heads = 8,
    dropout = 0.,
    cond_drop_prob = 0.25,  # for classifier free guidance
    ff_mult = 4,
    t5_name = 't5-large',
)

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

loss = parti(
    texts = texts,
    images = images,
    return_loss = True
)

loss.backward()

  4. Pre-encode text with t5_encode_text for efficient training

    main

    For large-scale training, avoid re-encoding text in every step. Use t5_encode_text from parti_pytorch.t5 to pre-compute text_token_embeds and text_mask. These can be stored and loaded via a dataloader to be passed directly into the Parti instance.

    from parti_pytorch.t5 import t5_encode_text
import torch

images = torch.randn(4, 3, 256, 256).cuda()

# Pre-encode text
text_token_embeds, text_mask = t5_encode_text([
    'a child screaming at finding a worm within a half-eaten apple',
    'lizard running across the desert on two feet',
    'waking up to a psychedelic landscape',
    'seashells sparkling in the shallow waters'
], name = 't5-large', output_device = images.device)

# Use pre-encoded tokens in training
loss = parti(
    text_token_embeds = text_token_embeds,
    text_mask = text_mask,
    images = images,
    return_loss = True
)

loss.backward()

  5. Generate images with Parti

    main

    Use the .generate() method to produce images from text prompts. You can control the conditioning scale for classifier-free guidance using the cond_scale parameter. Setting return_pil_images = True returns a List[PILImages] (256 x 256 RGB).

    images = parti.generate(
    texts = [
        'a whale breaching from afar',
        'young girl blowing out candles on her birthday cake',
        'fireworks with blue and green sparkles'
    ],
    cond_scale = 3., 
    return_pil_images = True
)