### Install parti-pytorch Source: https://github.com/lucidrains/parti-pytorch/blob/main/README.md Install the package via pip. ```bash $ pip install parti-pytorch ``` -------------------------------- ### Forward Pass with Classifier-Free Guidance Source: https://context7.com/lucidrains/parti-pytorch/llms.txt The `forward_with_cond_scale` method is used internally for classifier-free guidance during inference. This example shows the setup for such a pass. ```python import torch from parti_pytorch import Parti, VitVQGanVAE from parti_pytorch.t5 import t5_encode_text # Setup model (assuming pre-trained weights) vit_vae = VitVQGanVAE(dim=256, image_size=256, patch_size=16, num_layers=3).cuda() parti = Parti(vae=vit_vae, dim=512, depth=8, dim_head=64, heads=8, t5_name='t5-large') ``` -------------------------------- ### Initialize and Use Parti Source: https://github.com/lucidrains/parti-pytorch/blob/main/README.md Load a trained VAE, configure the Parti model, and perform training or generation. ```python import torch from parti_pytorch import Parti, VitVQGanVAE # first instantiate your ViT VQGan VAE # a VQGan VAE made of transformers 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() vit_vae.load_state_dict(torch.load(f'/path/to/vae.pt')) # you will want to load the exponentially moving averaged VAE # then you plugin the ViT VqGan VAE into your Parti as so parti = Parti( vae = vit_vae, # vit vqgan vae dim = 512, # model dimension depth = 8, # depth dim_head = 64, # attention head dimension heads = 8, # attention heads dropout = 0., # dropout cond_drop_prob = 0.25, # conditional dropout, for classifier free guidance ff_mult = 4, # feedforward expansion factor t5_name = 't5-large', # name of your T5 ) # ready your training text and images texts = [ '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' ] images = torch.randn(4, 3, 256, 256).cuda() # feed it into your parti instance, with return_loss set to True loss = parti( texts = texts, images = images, return_loss = True ) loss.backward() # do this for a long time on much data # then... 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) # conditioning scale for classifier free guidance # List[PILImages] (256 x 256 RGB) ``` -------------------------------- ### Initialize and Use Parti Model for Training and Generation Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Initialize the VAE and Parti model, then use it for training by computing loss with text-image pairs, or for generation from text prompts. Ensure VAE weights are loaded before initializing Parti. ```python import torch from parti_pytorch import Parti, VitVQGanVAE # Initialize the pre-trained VAE vit_vae = VitVQGanVAE( dim = 256, image_size = 256, patch_size = 16, num_layers = 3 ).cuda() # Load pre-trained VAE weights vit_vae.load_state_dict(torch.load('/path/to/vae.pt')) # Initialize Parti with the VAE parti = Parti( vae = vit_vae, # Pre-trained ViT VQ-GAN VAE dim = 512, # Transformer model dimension depth = 8, # Number of transformer layers dim_head = 64, # Attention head dimension heads = 8, # Number of attention heads dropout = 0., # Dropout rate cond_drop_prob = 0.25, # Conditional dropout for CFG ff_mult = 4, # Feedforward expansion factor t5_name = 't5-large', # T5 model for text encoding max_text_len = 128 # Maximum text sequence length ) # Training: compute loss with text-image pairs texts = [ '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' ] images = torch.randn(4, 3, 256, 256).cuda() loss = parti( texts = texts, images = images, return_loss = True ) loss.backward() print(f"Training loss: {loss.item()}") # Generation: create images from text prompts generated_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., # Classifier-free guidance scale filter_thres = 0.9, # Top-k filtering threshold temperature = 1., # Sampling temperature return_pil_images = True # Return PIL images ) # Save generated images for i, img in enumerate(generated_images): img.save(f'generated_{i}.png') ``` -------------------------------- ### Initialize and Use VQGanVAETrainer Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Initialize the VQGanVAETrainer for a complete training loop of the VitVQGanVAE. Configure dataset loading, gradient accumulation, mixed-precision, EMA, and checkpoint saving. After training, load the EMA model for best results. ```python from parti_pytorch import VitVQGanVAE, VQGanVAETrainer # Initialize the VAE model vit_vae = VitVQGanVAE( dim = 256, image_size = 256, patch_size = 16, num_layers = 3 ).cuda() # Create the trainer trainer = VQGanVAETrainer( vit_vae, folder = '/path/to/your/images', # Path to image folder num_train_steps = 100000, # Total training steps lr = 3e-4, # Learning rate batch_size = 4, # Batch size grad_accum_every = 8, # Gradient accumulation steps wd = 0., # Weight decay amp = True, # Enable mixed precision training save_results_every = 100, # Save sample reconstructions every N steps save_model_every = 1000, # Save model checkpoint every N steps results_folder = './results', # Output folder for checkpoints valid_frac = 0.05, # Fraction for validation split ema_beta = 0.995, # EMA decay rate ema_update_after_step = 500, # Start EMA after N steps ema_update_every = 10, # Update EMA every N steps apply_grad_penalty_every = 4 # Apply gradient penalty every N steps ) # Train the VAE (saves checkpoints to results_folder) trainer.train() # After training, load the EMA model for best results vit_vae.load_state_dict(torch.load('./results/vae.100000.ema.pt')) ``` -------------------------------- ### Initialize and Use VitVQGanVAE Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Initialize the Vision Transformer VQ-GAN Variational Autoencoder for image tokenization. Encode images to discrete tokens and decode them back. Supports reconstruction and discriminator losses for training. ```python import torch from parti_pytorch import VitVQGanVAE # Initialize the VitVQGanVAE vit_vae = VitVQGanVAE( dim = 256, # Model dimensions image_size = 256, # Target image size (must be square) channels = 3, # Number of image channels (RGB) layers = 4, # Number of transformer layers patch_size = 16, # Patch size for tokenization codebook_size = 65536, # Size of the quantization codebook lookup_free_quantization = True, # Use LFQ (recommended) l2_recon_loss = False, # Use L1 loss (False) or L2 loss (True) use_hinge_loss = True, # Use hinge loss for GAN training use_vgg_and_gan = True, # Enable perceptual and adversarial losses discr_layers = 4 # Number of discriminator layers ).cuda() # Encode an image to discrete tokens images = torch.randn(4, 3, 256, 256).cuda() fmap, indices, vq_loss = vit_vae.encode(images, return_indices_and_loss=True) print(f"Encoded feature map shape: {fmap.shape}") print(f"Token indices shape: {indices.shape}") # Decode tokens back to image reconstructed = vit_vae.decode(fmap) print(f"Reconstructed image shape: {reconstructed.shape}") # Forward pass with reconstruction loss (for training) loss = vit_vae(images, return_loss=True) loss.backward() print(f"VAE loss: {loss.item()}") # Forward pass with discriminator loss (for GAN training) discr_loss = vit_vae(images, return_discr_loss=True) discr_loss.backward() print(f"Discriminator loss: {discr_loss.item()}") # Get feature map from codebook indices (for generation) indices_2d = indices.view(4, 16, 16) # Reshape to 2D grid fmap_from_codes = vit_vae.get_fmap_from_codebook(indices_2d) decoded_from_codes = vit_vae.decode(fmap_from_codes) ``` -------------------------------- ### Train ViT VQGan VAE Source: https://github.com/lucidrains/parti-pytorch/blob/main/README.md Initialize and train the VQGan VAE model before using it with Parti. ```python 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() ``` -------------------------------- ### Generate Image Logits with Classifier-Free Guidance Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Encodes text using T5 and computes logits with varying conditioning scales to control text adherence. ```python # Encode text texts = ['a beautiful sunset over the ocean'] text_embeds, text_mask = t5_encode_text(texts, name='t5-large', output_device='cuda') # Get logits with classifier-free guidance # Higher cond_scale = stronger text conditioning (typical: 1.0 to 5.0) image_token_ids = torch.randint(0, 65536, (1, 100)).cuda() # Partial sequence # Without CFG (cond_scale=1) logits_no_cfg = parti.forward_with_cond_scale( text_token_embeds = text_embeds, text_mask = text_mask, image_token_ids = image_token_ids, cond_scale = 1.0 ) # With CFG (cond_scale=3) logits_with_cfg = parti.forward_with_cond_scale( text_token_embeds = text_embeds, text_mask = text_mask, image_token_ids = image_token_ids, cond_scale = 3.0 ) print(f"Logits shape: {logits_with_cfg.shape}") # (1, 101, codebook_size) ``` -------------------------------- ### Pre-encode Text for Scaling Source: https://github.com/lucidrains/parti-pytorch/blob/main/README.md Use T5 to pre-encode text into tokens and masks for more efficient training at scale. ```python from parti_pytorch.t5 import t5_encode_text images = torch.randn(4, 3, 256, 256).cuda() 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) # store somewhere, then load with the dataloader loss = parti( text_token_embeds = text_token_embeds, text_mask = text_mask, images = images, return_loss = True ) loss.backward() ``` -------------------------------- ### Create Lightweight VAE for Inference Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Uses copy_for_eval to remove discriminator and VGG networks from the VAE, reducing memory footprint for inference tasks. ```python import torch from parti_pytorch import VitVQGanVAE # Initialize full VAE with discriminator and VGG vit_vae = VitVQGanVAE( dim = 256, image_size = 256, patch_size = 16, num_layers = 3, use_vgg_and_gan = True ).cuda() # Create lightweight copy for inference vae_eval = vit_vae.copy_for_eval() # The eval copy doesn't have discriminator/VGG (saves memory) print(f"Original VAE has discriminator: {hasattr(vit_vae, 'discr')}") # True print(f"Eval VAE has discriminator: {hasattr(vae_eval, 'discr')}") # False # Use eval copy for encoding/decoding images = torch.randn(2, 3, 256, 256).cuda() with torch.no_grad(): fmap = vae_eval.encode(images, return_indices_and_loss=False) reconstructed = vae_eval.decode(fmap) ``` -------------------------------- ### Pre-encode Text for Efficient Training with T5 Source: https://context7.com/lucidrains/parti-pytorch/llms.txt Use `t5_encode_text` to pre-encode text descriptions with T5, which speeds up training by avoiding repeated encoding. This function returns text embeddings and a mask. Embeddings can be saved and loaded for later use. ```python import torch from parti_pytorch import Parti, VitVQGanVAE from parti_pytorch.t5 import t5_encode_text, get_encoded_dim # Check the embedding dimension for a T5 model dim = get_encoded_dim('t5-large') print(f"T5-large embedding dimension: {dim}") # 1024 # Pre-encode text descriptions texts = [ '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' ] device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') text_token_embeds, text_mask = t5_encode_text( texts, name = 't5-large', # T5 model name output_device = device # Target device for embeddings ) print(f"Text embeddings shape: {text_token_embeds.shape}") # (4, seq_len, 1024) print(f"Text mask shape: {text_mask.shape}") # (4, seq_len) # Save embeddings for later use in training torch.save({ 'embeds': text_token_embeds, 'mask': text_mask }, 'precomputed_text_embeds.pt') # Load and use pre-encoded embeddings for training saved = torch.load('precomputed_text_embeds.pt') images = torch.randn(4, 3, 256, 256).to(device) # Initialize Parti (assuming vit_vae is already loaded) vit_vae = VitVQGanVAE(dim=256, image_size=256, patch_size=16, num_layers=3).to(device) parti = Parti(vae=vit_vae, dim=512, depth=8, dim_head=64, heads=8, t5_name='t5-large') # Train with pre-encoded text embeddings (faster) loss = parti( text_token_embeds = saved['embeds'], text_mask = saved['mask'], images = images, return_loss = True ) loss.backward() print(f"Training loss with pre-encoded text: {loss.item()}") ``` === COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.