### Install and Run Oxford Flowers Example

Source: https://github.com/lucidrains/discrete-distribution-network/blob/main/README.md

Commands to install uv and execute the training script.

```shell
$ pip install uv
```

```shell
$ uv run train_oxford_flowers.py
```

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

### Install Discrete Distribution Network

Source: https://github.com/lucidrains/discrete-distribution-network/blob/main/README.md

Install the package via pip.

```bash
$ pip install discrete-distribution-network
```

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

### Trainer Class for Complete Training Pipeline

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Illustrates the setup of the Trainer class, which provides a comprehensive training pipeline with Accelerate integration, EMA, checkpointing, and automatic split-and-prune calls.

```python
import torch
from torch import nn
from torch.utils.data import Dataset
from discrete_distribution_network import DDN, Trainer
import torchvision.transforms as T
from datasets import load_dataset
```

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

### Configure and Run Trainer

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Sets up the Trainer with the DDN model and dataset, configuring training steps, learning rate, batch size, and checkpointing. The trainer handles the training loop, sampling, and checkpointing automatically when called.

```python
# Create trainer with full configuration
trainer = Trainer(
    ddn,
    dataset=dataset,
    num_train_steps=70_000,              # Total training steps
    learning_rate=3e-4,                  # AdamW learning rate
    weight_decay=1e-3,                   # L2 regularization
    batch_size=32,                       # Training batch size
    checkpoints_folder='./checkpoints',  # Save checkpoints here
    results_folder='./results',          # Save samples here
    save_results_every=100,              # Sample generation frequency
    checkpoint_every=1000,               # Checkpoint frequency
    num_samples=16,                      # Samples per generation (must be square)
    use_ema=True,                        # Use EMA for sampling
    max_grad_norm=0.5,                   # Gradient clipping
    accelerate_kwargs=dict(),            # Accelerate configuration
    ema_kwargs=dict()                    # EMA configuration
)

# Run training (automatically handles split_and_prune, checkpointing, sampling)
trainer()
```

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

### Patch-Based GuidedSampler Initialization and Usage

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Initializes and uses the GuidedSampler for patch-based processing. This allows for finer-grained control and per-patch code selection, suitable for high-resolution images.

```python
import torch
from discrete_distribution_network import GuidedSampler

# Patch-based guided sampler
sampler = GuidedSampler(
    dim=16,
    dim_query=3,
    codebook_size=10,
    patch_size=16,                       # Process 16x16 patches
    min_total_count_before_split_prune=1
)

features = torch.randn(10, 16, 32, 32)   # 32x32 features = 2x2 patches of 16x16
query_image = torch.randn(10, 3, 32, 32)

output, codes, commit_loss = sampler(features, query_image)
print(output.shape)   # torch.Size([10, 3, 32, 32])
print(codes.shape)    # torch.Size([10, 2, 2]) - one code per patch

# With residual connection (for multi-stage processing)
residual = torch.randn(10, 3, 32, 32)
output_with_res, codes, loss = sampler(features, query_image, residual=residual)

# Forward with patch codes
with torch.no_grad():
    out = sampler.forward_for_codes(features[:3], codes[:3])
    print(out.shape)  # torch.Size([3, 3, 32, 32])
```

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

### Initialize Dataset and DDN Model

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Initializes the custom dataset and the DDN model with specified dimensions and configuration. The `guided_sampler_kwargs` can be tuned for specific sampling behaviors.

```python
IMAGE_SIZE = 64
dataset = OxfordFlowersDataset(image_size=IMAGE_SIZE)

ddn = DDN(
    dim=64,
    image_size=IMAGE_SIZE,
    dropout=0.05,
    guided_sampler_kwargs=dict(
        crossover_top2_prob=0.1,
        straight_through_distance_logits=True,
        pre_network_activation=nn.SiLU()
    )
)
```

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

### Use GuidedSampler

Source: https://github.com/lucidrains/discrete-distribution-network/blob/main/README.md

Configure and use the GuidedSampler module, including the split-and-prune operation.

```python
import torch
from discrete_distribution_network import GuidedSampler

sampler = GuidedSampler(
    dim = 16,              # feature dimension
    dim_query = 3,         # the query image dimension
    codebook_size = 10,    # the codebook size K in the paper, which is K separate projections of the features, which is then measured distance wise to the query image guide
)

features = torch.randn(20, 16, 32, 32)
query_image = torch.randn(20, 3, 32, 32)

out, codes, commit_loss = sampler(features, query_image)

# (20, 3, 32, 32), (20,), ()

assert torch.allclose(sampler.forward_for_codes(features, codes), out, atol = 1e-5)

# after optimizer step, this needs to be called
# there is also a helper function by the same name that can take in your module and will invoke all of the guided samplers

sampler.split_and_prune_()
```

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

### Implement GuidedSampler Module

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Configures the GuidedSampler for feature projection and demonstrates a forward pass with query guidance.

```python
import torch
from torch import nn
from discrete_distribution_network import GuidedSampler

# Basic guided sampler for image features
sampler = GuidedSampler(
	dim=16,                              # Input feature dimension
	dim_query=3,                         # Query/output dimension (3 for RGB)
	codebook_size=10,                    # K projections in codebook
	chain_dropout_prob=0.05,             # Random code selection probability
	split_thres=2.0,                     # Threshold for splitting frequent codes
	prune_thres=0.5,                     # Threshold for pruning rare codes
	min_total_count_before_split_prune=100,  # Min samples before split/prune
	prenorm=True,                        # Apply RMS normalization to inputs
)

# Forward pass with query guidance
features = torch.randn(20, 16, 32, 32)      # Feature maps
query_image = torch.randn(20, 3, 32, 32)    # Guide images (e.g., downsampled targets)

output, codes, commit_loss = sampler(features, query_image)
# output: torch.Size([20, 3, 32, 32]) - selected projections
# codes: torch.Size([20,]) - selected code indices per batch
# commit_loss: scalar - MSE between output and query

# Backprop and update
commit_loss.backward()
optimizer.step()
sampler.split_and_prune_()  # Balance codebook after each step
```

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

### Initialize and Train DDN

Source: https://github.com/lucidrains/discrete-distribution-network/blob/main/README.md

Basic usage for initializing the DDN module, performing a forward pass, and sampling.

```python
import torch
from discrete_distribution_network import DDN

ddn = DDN(
    dim = 32,
    image_size = 256
)

images = torch.randn(2, 3, 256, 256)

loss = ddn(images)
loss.backward()

# after much training

sampled = ddn.sample(batch_size = 1)

assert sampled.shape == (1, 3, 256, 256)
```

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

### Initialize and Train DDN Model

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Configures the DDN model with hierarchical ResNet blocks and performs a training forward pass followed by codebook balancing.

```python
import torch
from discrete_distribution_network import DDN

# Initialize the DDN model
ddn = DDN(
	dim=32,                    # Base feature dimension
	dim_max=1024,              # Maximum feature dimension (caps scaling)
	image_size=256,            # Output image size (must be power of 2)
	channels=3,                # Image channels (RGB=3)
	codebook_size=10,          # Number of discrete codes per stage (K in paper)
	dropout=0.05,              # Dropout probability
	num_resnet_blocks=2,       # ResNet blocks per stage
	guided_sampler_kwargs=dict(
		crossover_top2_prob=0.1,              # Probability of crossover during split
		straight_through_distance_logits=True  # Use straight-through gradient estimation
	)
)

# Training forward pass - returns commit loss
images = torch.randn(4, 3, 256, 256)  # Batch of images
loss = ddn(images)
loss.backward()

# Call split_and_prune after each optimizer step to balance codebook
ddn.split_and_prune_()

# Sampling - generate new images
with torch.no_grad():
	# Random sampling with batch size
	sampled_images = ddn.sample(batch_size=4)
	print(sampled_images.shape)  # torch.Size([4, 3, 256, 256])

	# Deterministic sampling with specific codes
	# codes shape: (batch, num_stages) where num_stages = log2(image_size)
	codes = torch.randint(0, 10, (2, 8))  # 8 stages for 256x256
	sampled_with_codes = ddn.sample(codes=codes)

	# Get codes alongside samples for reproducibility
	sampled, returned_codes = ddn.sample(batch_size=2, return_codes=True)
	print(returned_codes.shape)  # torch.Size([2, 8])
```

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

### Codebook Balancing with split_and_prune_

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Demonstrates the usage of `split_and_prune_` for balancing codebook usage in GuidedSampler and DDN models. This function splits frequent codes and prunes rare ones to prevent mode collapse.

```python
import torch
from discrete_distribution_network import DDN, GuidedSampler, split_and_prune_

# For a single GuidedSampler
sampler = GuidedSampler(
    dim=16,
    dim_query=3,
    codebook_size=10,
    split_thres=2.0,                     # Split if count > 2.0/K of total
    prune_thres=0.5,                     # Prune if count < 0.5/K of total
    min_total_count_before_split_prune=100,
    crossover_top2_prob=0.1              # Optionally crossover top-2 during split
)

# Training loop with split_and_prune
for batch in dataloader:
    features, query = batch
    output, codes, loss = sampler(features, query)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()
    sampler.split_and_prune_()           # Call after each optimizer step

# For entire DDN model (calls split_and_prune_ on all GuidedSamplers)
ddn = DDN(dim=32, image_size=64)
ddn.split_and_prune_()                   # Balances all samplers in the network

# Or use the module-level helper on any parent module
model = nn.Sequential(
    sampler,
    nn.Conv2d(3, 16, 3, padding=1)
)
split_and_prune_(model)                  # Finds and updates all GuidedSamplers
```

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

### GuidedSampler for Non-Image Data with Custom Modules

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Configures and uses the GuidedSampler for arbitrary tensor data, including 1D feature vectors, by providing custom network and normalization modules.

```python
import torch
from torch import nn
from discrete_distribution_network import GuidedSampler

# Guided sampler for 1D feature vectors
sampler = GuidedSampler(
    dim=32,                              # Input feature dimension
    dim_query=16,                        # Output dimension
    codebook_size=10,
    network=nn.Linear(32, 16),           # Custom projection network
    norm_module=nn.RMSNorm(32),          # Custom normalization
    prenorm=True,
    min_total_count_before_split_prune=1
)

features = torch.randn(10, 32)           # 1D features
query = torch.randn(10, 16)              # Target outputs

output, codes, commit_loss = sampler(features, query)
print(output.shape)   # torch.Size([10, 16])
print(codes.shape)    # torch.Size([10,])

# After training
sampler.split_and_prune_()
out = sampler.forward_for_codes(features[:3], codes[:3])
print(out.shape)      # torch.Size([3, 16])
```

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

### Inference with Specific Codes using GuidedSampler

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Demonstrates deterministic generation by providing specific codes to the GuidedSampler for a batch of features. Supports single code for the entire batch or different codes per element.

```python
with torch.no_grad():
    # Single code for entire batch
    out_single = sampler.forward_for_codes(features, codes=torch.tensor(5))

    # Different code per batch element
    batch_codes = torch.tensor([3, 5, 2, 7, 1, 0, 9, 4, 6, 8, 2, 1, 0, 5, 7, 3, 8, 9, 4, 6])
    out_batch = sampler.forward_for_codes(features, codes=batch_codes)

    # Verify reproducibility
    assert torch.allclose(sampler.forward_for_codes(features, codes), output, atol=1e-5)
```

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

### Citation

Source: https://github.com/lucidrains/discrete-distribution-network/blob/main/README.md

BibTeX citation for the Discrete Distribution Networks paper.

```bibtex
@misc{yang2025discretedistributionnetworks,
    title   = {Discrete Distribution Networks}, 
    author  = {Lei Yang},
    year    = {2025},
    eprint  = {2401.00036},
    archivePrefix = {arXiv},
    primaryClass = {cs.CV},
    url     = {https://arxiv.org/abs/2401.00036}, 
}
```

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

### Define Custom Dataset

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Defines a custom PyTorch Dataset for the Oxford Flowers dataset, loading images and applying transformations. Ensure the dataset is compatible with PyTorch's Dataset API.

```python
class OxfordFlowersDataset(Dataset):
    def __init__(self, image_size):
        self.ds = load_dataset('nelorth/oxford-flowers')['train']
        self.transform = T.Compose([
            T.RandomHorizontalFlip(),
            T.RandomResizedCrop((image_size, image_size), (0.8, 1.0)),
            T.PILToTensor()
        ])

    def __len__(self):
        return len(self.ds)

    def __getitem__(self, idx):
        pil = self.ds[idx]['image']
        tensor = self.transform(pil)
        return tensor / 255.0
```

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

### Manual Trainer Operations

Source: https://context7.com/lucidrains/discrete-distribution-network/llms.txt

Performs manual operations on the trainer, including saving the current model checkpoint, loading a previously saved checkpoint, and generating samples. These operations allow for fine-grained control outside the automatic training loop.

```python
# Manual operations
trainer.save('custom_checkpoint.pt')     # Save checkpoint
trainer.load('custom_checkpoint.pt')     # Load checkpoint
trainer.sample('manual_sample.png')      # Generate and save samples
```

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