### Complete Training Pipeline Example

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

Illustrates a complete training pipeline for video generation, involving the VideoTokenizer, Gaia2 model, and ReconDiscriminator.

```python
import torch
import torch.optim as optim
from gaia2_pytorch import VideoTokenizer, Gaia2, ReconDiscriminator
```

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

### Install Gaia2-Pytorch

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

Install the package via pip.

```bash
$ pip install gaia2-pytorch
```

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

### Install Development Dependencies

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

Install dependencies required for running tests.

```bash
$ pip install '.[test]'
```

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

### Initialize Gaia2 components and optimizers

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

Sets up the VideoTokenizer, ReconDiscriminator, and Gaia2 model, along with their respective Adam optimizers.

```python
tokenizer = VideoTokenizer(
    dim=256,
    dim_latent=32,
    enc_depth=2,
    dec_depth=2
)

discriminator = ReconDiscriminator(dim=64, depth=2)

gaia2 = Gaia2(
    tokenizer=tokenizer,
    dim_latent=32,
    dim=256,
    depth=4,
    heads=8
)

# Optimizers
tokenizer_opt = optim.Adam(tokenizer.parameters(), lr=1e-4)
discriminator_opt = optim.Adam(discriminator.parameters(), lr=1e-4)
gaia2_opt = optim.Adam(gaia2.parameters(), lr=1e-4)
```

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

### Initialize and Use Gaia2 Components

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

Demonstrates tokenizing video data and generating outputs using the Gaia2 model.

```python
import torch
from gaia2_pytorch import VideoTokenizer, Gaia2

video = torch.randn(1, 3, 10, 16, 16)

tokenizer = VideoTokenizer()

loss = tokenizer(video)
loss.backward()

gaia2 = Gaia2(tokenizer)

loss = gaia2(video)
loss.backward()

generated = gaia2.generate((10, 16, 16))
assert generated.shape == video.shape
```

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

### Initialize and Use VideoTokenizer

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

Initialize the VideoTokenizer for encoding and decoding videos. Expects videos in shape (batch, channels, time, height, width). The training forward pass returns the total loss for backpropagation.

```python
import torch
from gaia2_pytorch import VideoTokenizer

# Initialize the video tokenizer
# Expects videos in shape: (batch, channels, time, height, width)
tokenizer = VideoTokenizer(
    channels=3,          # RGB video
    dim=512,             # Transformer dimension
    dim_latent=64,       # Latent space dimension (small for efficiency)
    enc_depth=2,         # Encoder transformer depth
    dec_depth=2,         # Decoder transformer depth
    latent_loss_weight=1.,
    lpips_loss_weight=1.
)

# Create sample video: batch=1, channels=3, frames=10, height=16, width=16
video = torch.randn(1, 3, 10, 16, 16)

# Training forward pass - returns total loss for backpropagation
loss = tokenizer(video)
loss.backward()

# Get loss breakdown: (recon_loss, lpips_loss, latent_loss, adv_gen_loss)
total_loss, breakdown = tokenizer(video, return_breakdown=True)
print(f"Reconstruction: {breakdown[0]:.4f}, LPIPS: {breakdown[1]:.4f}, Latent KL: {breakdown[2]:.4f}")

# Encode video to latent space
mean, var = tokenizer.encode(video)
print(f"Latent shape: {mean.shape}")  # (batch, time, height, width, dim_latent)

# Sample from latent distribution
sampled_latents = tokenizer.encode(video, return_sampled=True)

# Decode latents back to video
reconstructed_video = tokenizer.decode(sampled_latents)
print(f"Reconstructed shape: {reconstructed_video.shape}")  # Same as input video
```

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

### Gaia2 with Context Conditioning

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

Initializes Gaia2 to work directly with latents and supports cross-attention to external context tokens for conditional generation. Training and generation with context conditioning are demonstrated.

```python
import torch
from gaia2_pytorch import Gaia2

# Initialize Gaia2 without tokenizer (work directly with latents)
model = Gaia2(
    tokenizer=None,          # No tokenizer - use raw latents
    dim_latent=77,           # Latent dimension
    dim=512,
    depth=8,
    heads=8,
    dim_context=256          # Context token dimension
)

# Input latents: (batch, time, height, width, dim_latent)
latents = torch.randn(2, 8, 16, 16, 77)

# Context tokens for cross-attention (e.g., action embeddings, text features)
context = torch.randn(2, 32, 256)       # (batch, num_tokens, dim_context)
context_mask = torch.ones(2, 32).bool() # Mask for variable-length context

# Training with context conditioning
loss = model(
    latents,
    context=context,
    context_mask=context_mask
)
loss.backward()

# Forward pass without loss (get predicted flow)
pred_flow = model(
    latents,
    context=context,
    context_mask=context_mask,
    return_flow_loss=False
)
print(f"Predicted flow shape: {pred_flow.shape}")  # Same as input latents

# Generate with classifier-free guidance
generated = model.generate(
    video_shape=(8, 16, 16),
    batch_size=2,
    steps=16,
    cond_scale=3.0  # Higher = stronger conditioning adherence
)
```

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

### Initialize and Use Gaia2 Model

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

Initialize the Gaia2 generative world model with a tokenizer. It performs flow matching on video latents using a factorized space-time transformer. The training forward pass returns the flow matching loss.

```python
import torch
from gaia2_pytorch import VideoTokenizer, Gaia2

# Create tokenizer
tokenizer = VideoTokenizer()

# Initialize Gaia2 with tokenizer
gaia2 = Gaia2(
    tokenizer=tokenizer,
    dim_latent=64,           # Must match tokenizer's dim_latent
    dim=512,                 # Transformer hidden dimension
    depth=24,                # Transformer depth (24 in paper)
    heads=16,                # Attention heads
    dim_head=64,             # Dimension per head
    context_dropout_prob=0.25,  # For classifier-free guidance
    use_logit_norm_distr=True   # Bimodal time sampling from paper
)

# Sample video for training
video = torch.randn(1, 3, 10, 16, 16)

# Training forward pass - returns flow matching loss
loss = gaia2(video)
loss.backward()

# Generate new video
```

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

### Run Tests

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

Execute the test suite using pytest.

```bash
$ pytest tests
```

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

### Citations

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

BibTeX citations for the research papers related to the implementation.

```bibtex
@article{Russell2025GAIA2AC,
    title   = {GAIA-2: A Controllable Multi-View Generative World Model for Autonomous Driving},
    author  = {Lloyd Russell and Anthony Hu and Lorenzo Bertoni and George Fedoseev and Jamie Shotton and Elahe Arani and Gianluca Corrado},
    journal = {ArXiv},
    year    = {2025},
    volume  = {abs/2503.20523},
    url     = {https://api.semanticscholar.org/CorpusID:277321454}
}
```

```bibtex
@article{Rombach2021HighResolutionIS,
    title   = {High-Resolution Image Synthesis with Latent Diffusion Models},
    author  = {Robin Rombach and A. Blattmann and Dominik Lorenz and Patrick Esser and Bj{"o}rn Ommer},
    journal = {2022 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
    year    = {2021},
    pages   = {10674-10685},
    url     = {https://api.semanticscholar.org/CorpusID:245335280}
}
```

```bibtex
@article{Zhu2025FracConnectionsFE,
    title   = {Frac-Connections: Fractional Extension of Hyper-Connections},
    author  = {Defa Zhu and Hongzhi Huang and Jundong Zhou and Zihao Huang and Yutao Zeng and Banggu Wu and Qiyang Min and Xun Zhou},
    journal = {ArXiv},
    year    = {2025},
    volume  = {abs/2503.14125},
    url     = {https://api.semanticscholar.org/CorpusID:277104144}
}
```

```bibtex
@inproceedings{Huang2025TheGI,
    title   = {The GAN is dead; long live the GAN! A Modern GAN Baseline},
    author  = {Yiwen Huang and Aaron Gokaslan and Volodymyr Kuleshov and James Tompkin},
    year    = {2025},
    url     = {https://api.semanticscholar.org/CorpusID:275405495}
}
```

```bibtex
@inproceedings{Darcet2023VisionTN,
    title   = {Vision Transformers Need Registers},
    author  = {Timoth'ee Darcet and Maxime Oquab and Julien Mairal and Piotr Bojanowski},
    year    = {2023},
    url     = {https://api.semanticscholar.org/CorpusID:263134283}
}
```

```bibtex
@misc{chen2025deepcompressionautoencoderefficient,
    title   = {Deep Compression Autoencoder for Efficient High-Resolution Diffusion Models},
    author  = {Junyu Chen and Han Cai and Junsong Chen and Enze Xie and Shang Yang and Haotian Tang and Muyang Li and Yao Lu and Song Han},
    year    = {2025},
    eprint  = {2410.10733},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV},
    url     = {https://arxiv.org/abs/2410.10733},
}
```

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

### Initialize and Use ReconDiscriminator

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

Initialize the ReconDiscriminator for adversarial training of the video tokenizer. It uses hinge loss and supports gradient penalty for stable GAN training. The discriminator loss is returned when `return_discr_loss=True`.

```python
import torch
from gaia2_pytorch import VideoTokenizer, ReconDiscriminator

# Initialize tokenizer and discriminator
tokenizer = VideoTokenizer(enc_depth=1, dec_depth=1)
discriminator = ReconDiscriminator(
    dim=32,       # Base channel dimension
    channels=3,   # RGB
    depth=2       # Number of conv blocks
)

video = torch.randn(1, 3, 10, 16, 16)

# Train discriminator: pass both reconstructed and real videos
# Returns hinge discriminator loss
discr_loss = tokenizer(
    video,
    recon_discr=discriminator,
    return_discr_loss=True,
    apply_grad_penalty=True  # Adds gradient penalty for stability
)
discr_loss.backward()

# Train generator: pass only through tokenizer with discriminator for adversarial loss
gen_loss = tokenizer(video, recon_discr=discriminator)
gen_loss.backward()

# Get loss breakdown for monitoring
total_loss, breakdown = tokenizer(
    video,
    recon_discr=discriminator,
    return_breakdown=True
)
# breakdown = (recon_loss, lpips_loss, latent_loss, adv_gen_loss)
```

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

### Execute Gaia2 training loop

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

Performs a three-step training process: updating the discriminator, the tokenizer, and the Gaia2 flow matching model.

```python
# Training loop
for step in range(1000):
    # Sample batch of videos: (batch, channels, time, height, width)
    video = torch.randn(4, 3, 10, 16, 16)

    # Step 1: Train discriminator
    discriminator_opt.zero_grad()
    discr_loss = tokenizer(video, recon_discr=discriminator, return_discr_loss=True)
    discr_loss.backward()
    discriminator_opt.step()

    # Step 2: Train tokenizer (VAE + adversarial)
    tokenizer_opt.zero_grad()
    tokenizer_loss = tokenizer(video, recon_discr=discriminator)
    tokenizer_loss.backward()
    tokenizer_opt.step()

    # Step 3: Train Gaia2 (flow matching)
    gaia2_opt.zero_grad()
    gaia2_loss = gaia2(video)
    gaia2_loss.backward()
    gaia2_opt.step()

    if step % 100 == 0:
        print(f"Step {step}: D={discr_loss:.4f}, VAE={tokenizer_loss:.4f}, Flow={gaia2_loss:.4f}")
```

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

### Generate Videos with Gaia2

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

Generates videos with a specified latent shape, batch size, and number of ODE solver steps. Classifier-free guidance scale can be adjusted for stronger conditioning adherence.

```python
import torch
from gaia2_pytorch import Gaia2

# Initialize Gaia2 model
gai2 = Gaia2(
    dim_latent=77,
    dim=512,
    depth=8,
    heads=8
)

# Generate video
generated_video = gaia2.generate(
    video_shape=(10, 16, 16),
    batch_size=1,
    steps=16,        # ODE solver steps
    cond_scale=3.0   # Classifier-free guidance scale
)
print(f"Generated video shape: {generated_video.shape}")  # (1, 3, 10, 16, 16)
```

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

### Generate videos with Gaia2

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

Switches the model to evaluation mode and generates new video sequences using the trained Gaia2 model.

```python
# Generate new videos after training
gaia2.eval()
with torch.no_grad():
    generated_videos = gaia2.generate(
        video_shape=(10, 16, 16),
        batch_size=4,
        steps=32,
        cond_scale=2.0
    )
print(f"Generated {generated_videos.shape[0]} videos of shape {generated_videos.shape[1:]}")
```

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

### Residual Downsampling and Upsampling

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

Implements spatial and temporal downsampling/upsampling modules with learned residual connections. These modules are used for efficient high-resolution latent processing.

```python
import torch
from gaia2_pytorch.gaia2 import ResidualDownsample, ResidualUpsample

# Input latents: (batch, time, height, width, dim)
latents = torch.randn(1, 10, 16, 16, 32)

# Spatial downsampling (2x reduction in height and width)
spatial_down = ResidualDownsample(dim=32, space=True)
downsampled = spatial_down(latents)
print(f"Spatial downsample: {latents.shape} -> {downsampled.shape}")
# Output: (1, 10, 8, 8, 64) - spatial dims halved, channels doubled

# Spatial upsampling (2x increase in height and width)
spatial_up = ResidualUpsample(dim=64, space=True)
upsampled = spatial_up(downsampled)
print(f"Spatial upsample: {downsampled.shape} -> {upsampled.shape}")
# Output: (1, 10, 16, 16, 32) - back to original shape

# Temporal downsampling (2x reduction in time dimension)
temporal_down = ResidualDownsample(dim=32, time=True)
temp_downsampled = temporal_down(latents)
print(f"Temporal downsample: {latents.shape} -> {temp_downsampled.shape}")
# Output: (1, 5, 16, 16, 64) - time halved, channels doubled

# Temporal upsampling
temporal_up = ResidualUpsample(dim=64, time=True)
temp_upsampled = temporal_up(temp_downsampled)
print(f"Temporal upsample: {temp_downsampled.shape} -> {temp_upsampled.shape}")
# Output: (1, 10, 16, 16, 32) - back to original
```

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

### Transformer for Video Generation

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

Implements a Transformer with factorized space-time attention, optional cross-attention, hyper-connections, and register tokens. It processes input tokens and can utilize conditioning and context for generation.

```python
import torch
from gaia2_pytorch.gaia2 import Transformer

# Initialize transformer with space-time attention
transformer = Transformer(
    dim=256,
    depth=4,
    dim_head=64,
    heads=8,
    ff_expansion_factor=4.,
    has_time_attn=True,              # Enable temporal attention
    cross_attend=True,               # Enable cross-attention
    dim_cross_attended_tokens=128,   # Context token dimension
    accept_cond=True,                # Enable adaptive normalization
    num_register_tokens=16,          # Vision register tokens
    num_hyperconn_streams=1,         # Hyper-connection streams
    num_hyperconn_fracs=4            # Fractional hyper-connections
)

# Input tokens: (batch, time, height, width, dim)
tokens = torch.randn(2, 8, 16, 16, 256)

# Conditioning (e.g., time embedding)
cond = torch.randn(2, 256)

# Context for cross-attention
context = torch.randn(2, 32, 128)
context_mask = torch.ones(2, 32).bool()

# Forward pass
output = transformer(
    tokens,
    context=context,
    context_mask=context_mask,
    cond=cond
)
print(f"Transformer output shape: {output.shape}")  # (2, 8, 16, 16, 256)
```

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