### Complete Training Pipeline Example

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Shows a full two-stage training process: first for BQVAE, then for MaskBit transformer. Requires dataset and results folders to be set up.

```python
import torch
from maskbit_pytorch import BQVAE, MaskBit
from maskbit_pytorch.trainer import BQVAETrainer, MaskBitTrainer
from torchvision.utils import save_image

# Stage 1: Train Binary Quantization VAE
vae = BQVAE(
    image_size=128,
    dim=256,
    depth=3,
    entropy_loss_weight=0.1,
    gen_loss_weight=0.1
)

vae_trainer = BQVAETrainer(
    vae,
    folder='./dataset/images',
    num_train_steps=100000,
    batch_size=8,
    image_size=128,
    lr=1e-4,
    use_ema=True,
    results_folder='./vae_results'
)

vae_trainer()  # Train VAE

# Stage 2: Train MaskBit Transformer
# Load the trained VAE (or use EMA version)
vae_state = torch.load('./vae_results/vae.100000.ema.pt')
vae.load_state_dict(vae_state)

maskbit = MaskBit(
    vae,
    dim=512,
    bits_group_size=256,
    depth=8,
    heads=8,
    dim_head=64
)

maskbit_trainer = MaskBitTrainer(
    maskbit,
    folder='./dataset/images',
    num_train_steps=200000,
    batch_size=16,
    image_size=128,
    lr=1e-4,
    results_folder='./maskbit_results'
)

maskbit_trainer()  # Train MaskBit

# Stage 3: Generate Images
maskbit.eval()
with torch.no_grad():
    generated_images = maskbit.sample(
        batch_size=16,
        num_demasking_steps=24,
        temperature=0.9
    )
    save_image(generated_images, 'generated_samples.png', nrow=4)
```

--------------------------------

### DataLoader Usage Example

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Demonstrates how to use DataLoader for batching dataset items. Ensure dataset is properly defined and normalized.

```python
dataloader = DataLoader(
    dataset,
    batch_size=16,
    shuffle=True,
    num_workers=4
)

# Iterate over batches
for batch in dataloader:
    # batch shape: (16, 3, 64, 64)
    # Normalized to [0, 1] range
    pass
```

--------------------------------

### Initialize and Train BQVAE

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Configures the Binary Quantization VAE and demonstrates forward passes for loss computation, reconstruction, and bit extraction.

```python
import torch
from maskbit_pytorch import BQVAE

# Initialize the VAE with binary quantization
vae = BQVAE(
    image_size=64,           # Input image size (square images)
    dim=512,                 # Base channel dimension
    channels=3,              # Input channels (RGB)
    depth=2,                 # Number of encoder/decoder blocks
    proj_in_kernel_size=7,   # Initial projection kernel size
    entropy_loss_weight=0.1, # Weight for entropy regularization
    reg_loss_weight=1e-2,    # LeCam regularization weight
    gen_loss_weight=1e-1     # Generator adversarial loss weight
)

# Training forward pass - returns total loss
images = torch.randn(4, 3, 64, 64)  # Batch of 4 RGB images
loss = vae(images, return_loss=True)
loss.backward()

# Get reconstruction without loss computation
reconstructed = vae(images, return_loss=False)  # Shape: (4, 3, 64, 64)

# Get discriminator loss for separate discriminator training
discr_loss = vae(images, return_discr_loss=True)
discr_loss.backward()

# Get detailed losses for logging
total_loss, recon_images, (recon_loss, entropy_loss, gen_loss) = vae(
    images,
    return_loss=True,
    return_details=True
)

# Extract quantized binary bits for downstream use
bits = vae(images, return_quantized_bits=True)  # Shape: (4, dim, h, w)
bits_bool = vae(images, return_quantized_bits=True, return_bits_as_bool=True)
```

--------------------------------

### Initialize and Train MaskBit

Source: https://github.com/lucidrains/maskbit-pytorch/blob/main/README.md

Demonstrates the workflow for training the BQVAE autoencoder and the MaskBit transformer, followed by image sampling.

```python
import torch
from maskbit_pytorch import BQVAE, MaskBit

images = torch.randn(1, 3, 64, 64)

# train vae

vae = BQVAE(
    image_size = 64,
    dim = 512
)

loss = vae(images, return_loss = True)
loss.backward()

# train maskbit

maskbit = MaskBit(
    vae,
    dim = 512,
    bits_group_size = 512,
    depth = 2
)

loss = maskbit(images)
loss.backward()

# after much training

sampled_image = maskbit.sample() # (1, 3, 64, 64)
```

--------------------------------

### Initialize and Train MaskBit Transformer

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Sets up the MaskBit transformer using a pretrained BQVAE to perform iterative denoising on bit sequences.

```python
import torch
from maskbit_pytorch import BQVAE, MaskBit

# First, train or load a BQVAE
vae = BQVAE(image_size=64, dim=512)
# ... train vae ...

# Initialize MaskBit transformer
maskbit = MaskBit(
    vae,                          # Pretrained BQVAE (will be set to eval mode)
    dim=512,                      # Transformer hidden dimension
    bits_group_size=512,          # Bits per "token" for grouping
    depth=2,                      # Number of transformer layers
    bits_groups=2,                # Consecutive bits masked together
    dim_head=64,                  # Attention head dimension
    heads=8,                      # Number of attention heads
    train_frac_bits_flipped=0.05, # Fraction of bits randomly flipped during training
    encoder_kwargs=dict(          # Additional x-transformers Encoder kwargs
        ff_mult=4,
        attn_flash=True
    )
)

# Training forward pass
images = torch.randn(4, 3, 64, 64)
loss = maskbit(images)
loss.backward()
```

--------------------------------

### Load Image Dataset

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Utility for loading and preprocessing images from a directory for training.

```python
from maskbit_pytorch.trainer import ImageDataset
from torch.utils.data import DataLoader

# Create dataset from folder of images
dataset = ImageDataset(
    folder='/path/to/images',    # Root folder with images
    image_size=64,               # Target size (images resized and center-cropped)
    exts=['jpg', 'jpeg', 'png']  # Supported file extensions
)
# Output: "1234 training samples found at /path/to/images"
```

--------------------------------

### Train BQVAE Model

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

The BQVAETrainer handles the full training pipeline for BQVAE, including distributed training, EMA, and checkpointing.

```python
import torch
from maskbit_pytorch import BQVAE
from maskbit_pytorch.trainer import BQVAETrainer

# Initialize model
vae = BQVAE(image_size=64, dim=512, depth=2)

# Create trainer
trainer = BQVAETrainer(
    vae,
    folder='/path/to/images',       # Directory with training images
    num_train_steps=100000,         # Total training steps
    batch_size=16,                  # Batch size per GPU
    image_size=64,                  # Target image size
    lr=3e-4,                        # Learning rate
    grad_accum_every=1,             # Gradient accumulation steps
    max_grad_norm=1.0,              # Gradient clipping for VAE
    discr_max_grad_norm=1.0,        # Gradient clipping for discriminator
    save_results_every=100,         # Save sample reconstructions interval
    save_model_every=1000,          # Save checkpoint interval
    results_folder='./results',     # Output directory
    valid_frac=0.05,                # Validation split fraction
    use_ema=True,                   # Use exponential moving average
    ema_beta=0.995,                 # EMA decay rate
    ema_update_after_step=0,        # Start EMA after N steps
    ema_update_every=1,             # EMA update frequency
    accelerate_kwargs=dict(         # HuggingFace Accelerate config
        mixed_precision='fp16'
    )
)

# Run training
trainer()  # Trains for num_train_steps

# Save and load checkpoints manually
trainer.save('./checkpoint.pt')
trainer.load('./checkpoint.pt')
```

--------------------------------

### Generate Images with MaskBit

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Use the sample method to generate images from noise using a trained MaskBit model. Supports custom batch sizes, demasking steps, and bit representation formats.

```python
import torch
from maskbit_pytorch import BQVAE, MaskBit

vae = BQVAE(image_size=64, dim=512)
maskbit = MaskBit(vae, dim=512, bits_group_size=512, depth=2)
# ... after training ...

# Generate a single image
generated = maskbit.sample()  # Shape: (1, 3, 64, 64)

# Generate batch with custom parameters
generated = maskbit.sample(
    batch_size=8,              # Number of images to generate
    num_demasking_steps=18,    # Iterative refinement steps
    temperature=1.0            # Sampling temperature (higher = more random)
)

# Get both images and bit representations
images, bits = maskbit.sample(
    batch_size=4,
    return_bits=True           # Also return the final bit sequence
)
# images shape: (4, 3, 64, 64)
# bits shape: (4, num_bits) with values -1 or +1

# Get bits as boolean values
images, bits_bool = maskbit.sample(
    batch_size=4,
    return_bits=True,
    return_bits_as_bool=True   # Bits as True/False instead of -1/+1
)
```

--------------------------------

### BQVAE Class

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

The BQVAE class handles binary quantization of images, providing methods for training, reconstruction, and bit extraction.

```APIDOC
## BQVAE Class

### Description
Initializes the Binary Quantization Variational Autoencoder for image compression into binary representations.

### Parameters
- **image_size** (int) - Required - Input image size (square).
- **dim** (int) - Required - Base channel dimension.
- **channels** (int) - Optional - Input channels (default 3).
- **depth** (int) - Optional - Number of encoder/decoder blocks.
- **entropy_loss_weight** (float) - Optional - Weight for entropy regularization.
- **reg_loss_weight** (float) - Optional - LeCam regularization weight.
- **gen_loss_weight** (float) - Optional - Generator adversarial loss weight.

### Methods
- **__call__(images, return_loss=False, return_discr_loss=False, return_details=False, return_quantized_bits=False, return_bits_as_bool=False)**: Performs forward pass for training or inference.
```

--------------------------------

### Train MaskBit Model

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Train the MaskBit transformer after the BQVAE is prepared. Requires a pretrained VAE checkpoint.

```python
import torch
from maskbit_pytorch import BQVAE, MaskBit
from maskbit_pytorch.trainer import MaskBitTrainer

# Load pretrained VAE
vae = BQVAE(image_size=64, dim=512, depth=2)
vae.load_state_dict(torch.load('./vae_checkpoint.pt'))

# Initialize MaskBit
maskbit = MaskBit(
    vae,
    dim=512,
    bits_group_size=512,
    depth=6,
    heads=8
)

# Create trainer
trainer = MaskBitTrainer(
    maskbit,
    folder='/path/to/images',       # Same dataset as VAE training
    num_train_steps=200000,         # Training steps
    batch_size=32,                  # Batch size
    image_size=64,                  # Image size (must match VAE)
    lr=3e-4,                        # Learning rate
    grad_accum_every=2,             # Gradient accumulation
    max_grad_norm=1.0,              # Gradient clipping
    save_results_every=500,         # Logging interval
    save_model_every=5000,          # Checkpoint interval
    results_folder='./maskbit_results',
    valid_frac=0.05,
    accelerate_kwargs=dict(
        mixed_precision='fp16'
    )
)

# Run training
trainer()

# After training, generate images
maskbit.eval()
with torch.no_grad():
    samples = maskbit.sample(batch_size=16, num_demasking_steps=18)
```

--------------------------------

### BQVAE.decode_bits_to_images

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Method to reconstruct images from binary bit representations.

```APIDOC
## BQVAE.decode_bits_to_images

### Description
Decodes binary bit representations back into images. Supports float tensors (-1/+1) or boolean tensors.

### Parameters
- **bits** (Tensor) - Required - Binary bit representation (float or boolean).
```

--------------------------------

### Decode Binary Bits to Images

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

Converts binary bit representations back into images, supporting float tensors, boolean tensors, and flattened sequences.

```python
import torch
from maskbit_pytorch import BQVAE

vae = BQVAE(image_size=64, dim=512)

# Decode from float bits (-1 or +1 values)
bits_float = torch.randn(2, 512, 16, 16).sign()  # Binary values as -1 or +1
images = vae.decode_bits_to_images(bits_float)  # Shape: (2, 3, 64, 64)

# Decode from boolean bits
bits_bool = torch.rand(2, 512, 16, 16) > 0.5
images = vae.decode_bits_to_images(bits_bool)  # Automatically converts to float

# Decode from flattened sequence
bits_flat = torch.randn(2, 512 * 16 * 16).sign()
images = vae.decode_bits_to_images(bits_flat)  # Reshapes automatically
```

--------------------------------

### MaskBit Citation

Source: https://github.com/lucidrains/maskbit-pytorch/blob/main/README.md

BibTeX citation for the MaskBit research paper.

```bibtex
@inproceedings{Weber2024MaskBitEI,
    title   = {MaskBit: Embedding-free Image Generation via Bit Tokens},
    author  = {Mark Weber and Lijun Yu and Qihang Yu and Xueqing Deng and Xiaohui Shen and Daniel Cremers and Liang-Chieh Chen},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:272832013}
}
```

--------------------------------

### MaskBit Class

Source: https://context7.com/lucidrains/maskbit-pytorch/llms.txt

The MaskBit class implements the transformer architecture for non-autoregressive mask decoding.

```APIDOC
## MaskBit Class

### Description
Initializes the transformer model that learns to predict masked bits through iterative denoising.

### Parameters
- **vae** (BQVAE) - Required - Pretrained BQVAE instance.
- **dim** (int) - Required - Transformer hidden dimension.
- **bits_group_size** (int) - Required - Bits per token.
- **depth** (int) - Required - Number of transformer layers.
- **bits_groups** (int) - Required - Consecutive bits masked together.
- **train_frac_bits_flipped** (float) - Optional - Fraction of bits flipped during training.

### Methods
- **__call__(images)**: Performs forward pass for training the transformer.
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.