### Install Examples for Rectified Flow PyTorch

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

Install the library with example dependencies to run the provided examples.

```bash
$ pip install .[examples]
```

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

### Run Oxford Flowers Example

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

Execute the training script for the Oxford Flowers dataset after installation.

```bash
$ python train_oxford.py
```

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

### Basic ImageDataset Setup

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initialize ImageDataset with a folder path and image size.

```python
from rectified_flow_pytorch import ImageDataset

dataset = ImageDataset(
    folder='./path/to/images',
    image_size=256
)
```

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

### Install Rectified Flow PyTorch

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

Install the library using pip. This is the first step before using any of the functionalities.

```bash
pip install rectified-flow-pytorch
```

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

### Basic Trainer Setup

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initializes the Trainer with a RectifiedFlow model, dataset, and training parameters. Ensure the dataset folder and image size are correctly specified.

```python
from rectified_flow_pytorch import Trainer, RectifiedFlow, ImageDataset

trainer = Trainer(
    RectifiedFlow(Unet(dim=64)),
    dataset=ImageDataset(folder='./images', image_size=256),
    num_train_steps=100_000,
    batch_size=32,
    learning_rate=3e-4
)

trainer()
```

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

### Reference-Guided Sampling (Follow the Mean)

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Guide generation using a reference image by providing `ref_image` and adjusting `ref_image_guide_strength` and `ref_image_max_time_guide`.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

model = Unet(dim=64)
rectified_flow = RectifiedFlow(model)

# Reference image to guide generation
reference = torch.randn(1, 3, 256, 256)

samples = rectified_flow.sample(
    batch_size=8,
    steps=32,
    ref_image=reference,
    ref_image_guide_strength=0.1,  # Strength of guidance
    ref_image_max_time_guide=0.85  # When to stop guiding
)
```

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

### Initialize and Run Trainer

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Instantiate the Trainer with a RectifiedFlow model and dataset, then start the training loop. Ensure dataset path and output folders are correctly specified.

```python
from rectified_flow_pytorch import Trainer, RectifiedFlow, ImageDataset, Unet

model = Unet(dim=64)
rectified_flow = RectifiedFlow(model)

dataset = ImageDataset(
    folder='./path/to/images',
    image_size=256
)

trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_train_steps=100_000,
    batch_size=32,
    results_folder='./results',
    checkpoints_folder='./checkpoints'
)

trainer()
```

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

### Train and Sample with RectifiedFlow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/README.md

Minimal example demonstrating how to initialize a Unet model, wrap it with RectifiedFlow, load an ImageDataset, set up a Trainer, and then train and sample from the model. Ensure you have an 'images' folder with training data.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet, Trainer, ImageDataset

# Create model
model = Unet(dim=64)
rectified_flow = RectifiedFlow(model)

# Load dataset
dataset = ImageDataset(folder='./images', image_size=256)

# Train
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_train_steps=100_000,
    batch_size=32
)
trainer()

# Sample
samples = rectified_flow.sample(batch_size=8, steps=32)
```

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

### Use a Smaller Model Architecture

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

For memory-constrained environments, consider using a model with fewer dimensions or layers. This example shows initializing a smaller Unet model.

```python
model = Unet(dim=32, dim_mults=(1, 2, 4))  # Smaller
```

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

### Example Usage of LossBreakdown

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/types.md

Demonstrates how to access individual loss components from the LossBreakdown named tuple after a forward pass.

```python
loss, breakdown = rectified_flow(data, return_loss_breakdown=True)
print(breakdown.total, breakdown.main, breakdown.data_match, breakdown.velocity_match)
```

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

### Use Perfect Square for num_samples

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

The `num_samples` parameter must be a perfect square (e.g., 4, 9, 16). This example shows initializing the Trainer with a valid `num_samples` value.

```python
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_samples=16  # 4x4 grid
)
```

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

### Generate Samples

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Use the `sample` method to generate data samples from the model. This method supports basic sampling, conditional generation, and guided sampling using a reference image. Specify `steps` for ODE solver precision.

```python
# Basic sampling
samples = rectified_flow.sample(batch_size=8, steps=32)
```

```python
# With conditioning
samples = rectified_flow.sample(
    batch_size=8, 
    steps=32,
    cond=class_labels
)
```

```python
# With reference guidance
samples = rectified_flow.sample(
    batch_size=8,
    steps=32,
    ref_image=ref_image.unsqueeze(0),
    ref_image_guide_strength=0.15
)
```

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

### One-Step Generation with MeanFlow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Use `MeanFlow` for one-step generation. The model should accept `(x, t, s)` where `s` is the integral start time. `use_adaptive_loss_weight` and `prob_default_flow_obj` can be configured.

```python
from rectified_flow_pytorch import MeanFlow
from torch.nn import Module

class SimpleModel(Module):
    def forward(self, x, t, s):
        # Mean flow expects (x, t, s) where s is integral start time
        return x * 0.1

mean_flow = MeanFlow(
    SimpleModel(),
    use_adaptive_loss_weight=True,
    prob_default_flow_obj=0.5
)

# One-step sampling (default)
samples = mean_flow.sample(batch_size=8)

# Multi-step sampling
samples = mean_flow.sample(batch_size=8, steps=8)
```

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

### Unet Forward Pass Example

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Demonstrates a typical forward pass of the Unet model, including optional conditioning. Ensure `accept_cond` and `dim_cond` match the model's initialization if using conditioning.

```python
model = Unet(dim=64, accept_cond=True, dim_cond=128)

x = torch.randn(8, 3, 256, 256)
times = torch.rand(8)
cond = torch.randn(8, 128)

output = model(x, times=times, cond=cond)  # (8, 3, 256, 256)
```

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

### Instantiate SelfFlow with FiT Model

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

SelfFlow requires a FiT (Flow Matching Transformer) instance as its primary model. This example demonstrates initializing SelfFlow with a FiT model.

```python
from rectified_flow_pytorch import SelfFlow, FiT

# Use FiT architecture
fit = FiT(dim=768, patch_size=16)
self_flow = SelfFlow(fit)
```

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

### Utilize Reflow for Model Refinement

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

This example shows how to use the Reflow mechanism for refining an existing Rectified Flow model. It covers both direct usage of the `Reflow` class and the `ReflowTrainer` for the refinement process, followed by sampling from the refined model.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet, Reflow, ReflowTrainer

# Train initial model
model = Unet(dim=64)
rectified_flow = RectifiedFlow(model)
# ... training code ...

# Apply reflow
reflow = Reflow(rectified_flow, batch_size=32)

# Or use trainer
reflow_trainer = ReflowTrainer(
    rectified_flow,
    num_train_steps=50_000,
    batch_size=32
)
reflow_trainer()

# Sample from refined model
samples = reflow.sample(batch_size=8)
```

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

### Multi-GPU Training Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Configure Rectified Flow for multi-GPU training using Accelerate. This setup automatically handles distribution and allows for mixed precision and gradient accumulation.

```python
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    accelerate_kwargs=dict(
        mixed_precision='fp16',
        gradient_accumulation_steps=2
    )
)
```

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

### Truncated Decay Function Example

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/types.md

An example implementation of a `DecayFn` which maps time gamma in [0, 1] to a decay value in [0, 1]. This specific function provides a constant decay up to a threshold `a`.

```python
DecayFn = Callable[[Tensor], Tensor]

def truncated_decay(gamma: Tensor, a: float = 0.8) -> Tensor:
    return torch.where(
        gamma <= a,
        torch.ones_like(gamma),
        (1. - gamma) / (1. - a)
    )
```

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

### Get Number of Images in Dataset

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Returns the total count of images found and loaded by the ImageDataset.

```python
def __len__(self) -> int:
    Returns number of images found.
```

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

### Trainer Sample Method

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Generates sample images and saves them to a specified file. The `fname` parameter indicates the path where the sample grid will be saved.

```python
def sample(self, fname: str) -> Tensor:
    # Generate samples and save to file.
    # Parameters:
    # - fname (str): Path to save the sample grid
    # Returns: Sampled images
```

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

### is_main property

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

A property that returns True if the current process is the main process in a distributed training setup.

```APIDOC
## is_main property

### Description
A property that returns True if the current process is the main process in a distributed training setup.

### Returns

- bool - True if this is the main process, False otherwise.
```

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

### UAFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initialize UAFlow (Uncertainty-Guided) with guidance scales and beta for NLL.

```python
from rectified_flow_pytorch import UAFlow

ua_flow = UAFlow(
    model,
    times_cond_kwarg='times',
    ucg_scale=1.0,        # Uncertainty guidance
    cfg_scale=1.0,        # Classifier-free guidance
    beta_nll=1.0
)
```

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

### Ensure Model is Float32

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

If the error persists, ensure the model itself is using float32. This example converts the model to float32.

```python
# Or ensure model is float32
model = model.float()
```

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

### ImageDataset Methods

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Provides methods to get the total number of images and retrieve individual images as normalized tensors.

```APIDOC
## ImageDataset Methods

### __len__

Returns the total number of images found in the dataset.

```python
def __len__(self) -> int:
```

### __getitem__

Returns a single image as a normalized tensor (0-1 range).

```python
def __getitem__(self, index: int) -> Tensor:
```

**Returns:** Tensor of shape (channels, height, width)
```

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

### sample

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Generates sample images and saves them to a specified file path.

```APIDOC
## sample

### Description
Generates sample images and saves them to a specified file path.

### Method
`sample(fname: str)`

### Parameters

- **fname** (str) - Required - Path to save the sample grid

### Returns

- Sampled images (Tensor)
```

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

### Typical Training Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Set up a standard training configuration for Rectified Flow. Ensure the model, loss function, and trainer are correctly initialized. 'use_consistency=False' is recommended for longer training durations.

```python
model = Unet(dim=64, dim_mults=(1, 2, 4, 8))

rectified_flow = RectifiedFlow(
    model,
    predict='flow',
    loss_fn='pseudo_huber',
    use_consistency=False,  # Use for longer training
    odeint_kwargs=dict(atol=1e-5, rtol=1e-5, method='midpoint')
)

dataset = ImageDataset('./images', image_size=256)

trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_train_steps=100_000,
    batch_size=32,
    learning_rate=3e-4,
    grad_accum_every=1,
    max_grad_norm=0.5
)

trainer()
```

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

### Trainer is_main Property

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

A property that returns True if the current process is the main process in a distributed training setup, and False otherwise.

```python
@property
def is_main(self) -> bool:
    # Returns True if this is the main process (for distributed training).
```

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

### UAFlow Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize UAFlow for unconstrained alignment flow with uncertainty and classifier-free guidance. Demonstrates sample generation.

```python
from rectified_flow_pytorch import UAFlow

ua_flow = UAFlow(
    model,
    ucg_scale=1.0,
    cfg_scale=1.0,
    beta_nll=1.0
)

samples = ua_flow.sample(steps=16, batch_size=8)
```

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

### Monitor Loss Breakdown

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Set `return_loss_breakdown=True` to get detailed loss components during training. Useful for debugging and understanding loss contributions.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet, Trainer

rectified_flow = RectifiedFlow(
    Unet(dim=64),
    use_consistency=True
)

data = torch.randn(8, 3, 256, 256)

# Get loss components
loss, breakdown = rectified_flow(data, return_loss_breakdown=True)

print(f"Total Loss: {breakdown.total:.4f}")
print(f"Main Loss: {breakdown.main:.4f}")
print(f"Data Match Loss: {breakdown.data_match:.4f}")
print(f"Velocity Match Loss: {breakdown.velocity_match:.4f}")
```

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

### Get Unet Downsampling Factor

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Retrieves the total downsampling factor of the Unet model, which determines the minimum divisibility requirement for input dimensions.

```python
model = Unet(dim=64, dim_mults=(1,2,4,8))
print(model.downsample_factor)  # 16
# Input dimensions must be divisible by 16
```

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

### Package Entry Point Exports

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/module-exports.md

Lists all public exports from the package's __init__.py file, including core classes and various flow variants.

```python
from rectified_flow_pytorch import (
    # Core classes
    RectifiedFlow,
    Unet,
    Trainer,
    ImageDataset,
    
    # Flow variants
    NanoFlow,
    MeanFlow,
    Reflow,
    ReflowTrainer,
    SplitMeanFlow,
    EquilibriumMatching,
    LsdFlow,
    FiT,
    SelfFlow,
    UAFlow,
    LapFlow,
    LapFlowDiT,
)
```

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

### sample

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Generates samples from the model using ODE integration. Supports various sampling configurations including temperature and reference image guidance.

```APIDOC
## sample

### Description
Generates samples from the model using ODE integration. Supports various sampling configurations including temperature and reference image guidance.

### Method Signature
@torch.no_grad()
def sample(
    batch_size: int = 1,
    steps: int = 16,
    noise: Tensor | None = None,
    data_shape: tuple[int, ...] | None = None,
    temperature: float = 1.,
    use_ema: bool = False,
    ref_image: Tensor | None = None,
    ref_image_guide_strength: float = 0.1,
    ref_image_max_time_guide: float = 0.85,
    **model_kwargs
) -> Tensor

### Parameters
- `batch_size` (int): Number of samples to generate
- `steps` (int): Number of ODE solver steps
- `noise` (Tensor | None): Starting noise. If None, samples from N(0,I).
- `data_shape` (tuple | None): Output shape (e.g., (3, 256, 256)). Must be set or known from training.
- `temperature` (float): Temperature for sampling when mean_variance_net=True
- `use_ema` (bool): Use EMA model if available (requires use_consistency=True)
- `ref_image` (Tensor | None): Reference image for guided generation (Follow the Mean paper)
- `ref_image_guide_strength` (float): Strength of reference guidance
- `ref_image_max_time_guide` (float): Maximum time value for applying guidance
- `**model_kwargs`: Model conditioning kwargs (e.g., cond, image_cond)

### Returns
- Sampled data tensor, shape (batch_size, *data_shape)

### Example
```python
# Basic sampling
samples = rectified_flow.sample(batch_size=8, steps=32)

# With conditioning
samples = rectified_flow.sample(
    batch_size=8, 
    steps=32,
    cond=class_labels
)

# With reference guidance
samples = rectified_flow.sample(
    batch_size=8,
    steps=32,
    ref_image=ref_image.unsqueeze(0),
    ref_image_guide_strength=0.15
)
```
```

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

### Initialize ImageDataset

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Instantiate the ImageDataset by specifying the folder containing images, desired image size, accepted file extensions, and horizontal flip augmentation.

```python
from rectified_flow_pytorch import ImageDataset

dataset = ImageDataset(
    folder='./path/to/images',
    image_size=256,
    exts=['jpg', 'png'],
    augment_horizontal_flip=True
)
```

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

### Document Statistics Table

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/README.md

This table provides a statistical breakdown of the documentation, including the number of lines, sections, and examples for each document and the total across all documents.

```markdown
| Document | Lines | Sections | Examples |
|----------|-------|----------|----------|
| rectified-flow-reference.md | 1,200+ | 15 | 40+ |
| types.md | 600+ | 20 | 15+ |
| configuration.md | 800+ | 18 | 50+ |
| errors.md | 500+ | 15 | 25+ |
| api-usage.md | 700+ | 20 | 45+ |
| module-exports.md | 400+ | 12 | 10+ |
| **Total** | **4,200+** | **100** | **185+** |
```

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

### Configure Multi-GPU Training with Accelerate

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Set up `accelerate_kwargs` for mixed precision and device placement. Note the effective batch size calculation.

```python
from rectified_flow_pytorch import Trainer

trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    accelerate_kwargs=dict(
        mixed_precision='fp16',
        cpu=False,
        device_placement=True,
        gradient_accumulation_steps=2
    ),
    batch_size=16,
    grad_accum_every=4  # Effective batch size: 256
)

trainer()
```

```bash
accelerate launch train.py
# or with multi-GPU
accelerate launch --multi_gpu --num_processes 4 train.py
```

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

### Get Single Image from Dataset

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Retrieves a single image from the dataset at the specified index, returning it as a normalized tensor with values ranging from 0 to 1.

```python
def __getitem__(self, index: int) -> Tensor:
    Returns a single image as normalized tensor (0-1 range).

    **Returns:** Tensor of shape (channels, height, width)
```

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

### EquilibriumMatching Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Set up EquilibriumMatching with a decay function, multiplier, and optimizer for sampling.

```python
from rectified_flow_pytorch import EquilibriumMatching

eq_matching = EquilibriumMatching(
    model,
    decay_fn=truncated_decay,
    decay_kwargs={'a': 0.8},
    lambda_multiplier=4.0,
    sample_optim_klass=torch.optim.SGD,
    sample_optim_kwargs=dict(lr=0.003, momentum=0.35, nesterov=True)
)
```

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

### Trainer Class Initialization

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/README.md

Set up the Trainer for training generative models. Key parameters include the RectifiedFlow model, dataset, training steps, batch size, learning rate, and checkpointing frequency. EMA and gradient accumulation are supported.

```python
Trainer(
    rectified_flow,
    dataset,
    num_train_steps=70_000,
    batch_size=16,
    learning_rate=3e-4,
    results_folder='./results',
    checkpoints_folder='./checkpoints',
    save_results_every=100,
    checkpoint_every=1000,
    use_ema=True,
    grad_accum_every=1,
    max_grad_norm=0.5,
    # ... plus options for optimization and distribution
)
```

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

### Use Consistency Flow Matching without Reflow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

If `use_consistency=True` is set for RectifiedFlow, do not use Reflow. This example shows training directly with consistency flow matching.

```python
# Option 1: Use consistency FM instead of reflow (preferred)
rectified_flow = RectifiedFlow(use_consistency=True)
# Train directly, don't use Reflow
```

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

### EquilibriumMatching Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize EquilibriumMatching for energy-based generative modeling using gradient descent. Demonstrates loss computation and iterative sampling.

```python
from rectified_flow_pytorch import EquilibriumMatching

eq_matching = EquilibriumMatching(
    model,
    decay_fn=truncated_decay,
    sample_optim_klass=torch.optim.SGD
)

loss = eq_matching(data)
samples = eq_matching.sample(steps=100)  # Iterative optimization
```

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

### Train a Rectified Flow Model from Scratch

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

This snippet demonstrates the complete workflow for training a Rectified Flow model from scratch. It includes defining the Unet model, wrapping it with RectifiedFlow, setting up an ImageDataset, configuring the Trainer, and initiating the training process.

```python
import torch
from rectified_flow_pytorch import (
    RectifiedFlow, Unet, Trainer, ImageDataset
)

# Define model
model = Unet(
    dim=64,
    dim_mults=(1, 2, 4, 8),
    channels=3
)

# Wrap in RectifiedFlow
rectified_flow = RectifiedFlow(
    model,
    predict='flow',
    loss_fn='mse',
    odeint_kwargs=dict(atol=1e-5, rtol=1e-5, method='midpoint')
)

# Load dataset
dataset = ImageDataset(
    folder='./path/to/images',
    image_size=256,
    augment_horizontal_flip=True
)

# Configure trainer
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_train_steps=100_000,
    batch_size=32,
    learning_rate=3e-4,
    results_folder='./results',
    checkpoints_folder='./checkpoints',
    save_results_every=500,
    checkpoint_every=5000
)

# Run training
trainer()
```

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

### Check Checkpoint Directory Before Loading

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

Prevent `RuntimeError: Checkpoint directory not found` by verifying the checkpoint file's existence before attempting to load it. This example uses `Path.exists()`.

```python
from pathlib import Path

checkpoint_path = Path('./checkpoints/checkpoint.1000.pt')

if checkpoint_path.exists():
    trainer.load('checkpoint.1000.pt')
else:
    print("Checkpoint not found, starting from scratch")
```

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

### SplitMeanFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initialize SplitMeanFlow with adaptive loss weighting and default flow object probability.

```python
from rectified_flow_pytorch import SplitMeanFlow

split_flow = SplitMeanFlow(
    model,
    use_adaptive_loss_weight=True,
    prob_default_flow_obj=0.5,
    accept_cond=False
)
```

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

### Verify Dataset Folder Exists

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

Ensure the dataset folder specified in `ImageDataset` actually exists to prevent `FileNotFoundError`. This example checks for folder existence before creating the dataset.

```python
from pathlib import Path

folder = Path('./path/to/images')
assert folder.exists(), f"Folder {folder} not found"

dataset = ImageDataset(folder=folder, image_size=256)
```

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

### Save and Load Checkpoints

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Demonstrates manual saving and loading of model and optimizer states using PyTorch's torch.save and torch.load. The trainer also offers automatic saving capabilities.

```python
# Trainer saves automatically
trainer.save('model.pt')  # Manually save
trainer.load('model.pt')  # Load checkpoint

# Or manually
state = {
    'model': rectified_flow.state_dict(),
    'optimizer': optimizer.state_dict(),
    'step': current_step
}
torch.save(state, 'checkpoint.pt')

# Load
checkpoint = torch.load('checkpoint.pt')
rectified_flow.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])
```

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

### LapFlowDiT Initialization

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize LapFlowDiT, a Diffusion Transformer architecture optimized for Laplacian multiscale flow matching.

```python
from rectified_flow_pytorch import LapFlow, LapFlowDiT

# LapFlowDiT is a transformer architecture optimized for laplacian flows
lap_dit = LapFlowDiT(dim=768, depth=12)
```

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

### Use Reflow without Consistency Flow Matching

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

If planning to use Reflow, ensure `use_consistency=False` is set for RectifiedFlow. This example shows training first, then applying Reflow.

```python
# Option 2: Don't use consistency if planning to reflow
rectified_flow = RectifiedFlow(use_consistency=False)
# Train, then use Reflow
reflow = Reflow(rectified_flow)
```

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

### SelfFlow Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize SelfFlow for self-supervised flow matching using a teacher-student architecture. It takes a FiT model and supports returning loss breakdown.

```python
from rectified_flow_pytorch import SelfFlow, FiT

fit = FiT(dim=768, patch_size=16)
self_flow = SelfFlow(
    fit,
    mask_ratio=0.5,
    repr_loss_weight=1.0
)

loss, breakdown = self_flow(data, return_loss_breakdown=True)
samples = self_flow.sample(steps=16)
```

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

### Batch Processing for Large Datasets

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Provides a function to process large datasets by iterating through images in batches. This example includes image loading, transformation, and saving generated samples.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

rectified_flow = RectifiedFlow(Unet(dim=64))
rectified_flow.eval()

def process_dataset(input_folder, output_folder, batch_size=8):
    from pathlib import Path
    from PIL import Image
    import torchvision.transforms as T
    
    output = Path(output_folder)
    output.mkdir(exist_ok=True)
    
    images = list(Path(input_folder).glob('*.jpg'))
    transform = T.Compose([
        T.Resize(256),
        T.CenterCrop(256),
        T.ToTensor()
    ])
    
    for i in range(0, len(images), batch_size):
        batch_images = images[i:i+batch_size]
        batch = torch.stack([
            transform(Image.open(img)) for img in batch_images
        ])
        
        with torch.no_grad():
            samples = rectified_flow.sample(
                batch_size=batch.shape[0],
                data_shape=(3, 256, 256),
                steps=32
            )
        
        for j, (sample, img_path) in enumerate(zip(samples, batch_images)):
            out_path = output / f'generated_{img_path.stem}.png'
            T.ToPILImage()(sample.cpu()).save(out_path)

process_dataset('./input', './output')
```

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

### Checkpointing and Results Strategy

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Sets up a strategy for saving model checkpoints and generating sample results at specified intervals. Ensure result and checkpoint folders exist.

```python
trainer = Trainer(
    model,
    dataset=dataset,
    checkpoint_every=1000,      # Save every N steps
    save_results_every=100,     # Generate samples every N steps
    num_samples=16,             # Must be perfect square
    results_folder='./results',
    checkpoints_folder='./checkpoints',
    clear_results_folder=True   # Remove old results on init
)
```

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

### Instantiate SelfFlow with EMA Model

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

When using SelfFlow, the teacher model must be an instance of EMA. This example shows creating a SelfFlow instance, which automatically uses EMA by default.

```python
from ema_pytorch import EMA

fit = FiT(dim=768)

# Default: uses EMA automatically
self_flow = SelfFlow(fit)
```

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

### Manual Training Loop with Rectified Flow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Demonstrates a basic training loop for the Rectified Flow model. Ensure to load data and handle optimizer steps correctly.

```python
import torch
from torch.optim import Adam
from rectified_flow_pytorch import RectifiedFlow, Unet

model = Unet(dim=64)
rectified_flow = RectifiedFlow(model)
optimizer = Adam(model.parameters(), lr=3e-4)

for step in range(10_000):
    # Load batch
    batch = torch.randn(32, 3, 256, 256)
    
    # Compute loss
    loss = rectified_flow(batch)
    
    # Backward
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optimizer.step()
    
    # Update EMA if using consistency
    if hasattr(rectified_flow, 'post_training_step_update'):
        rectified_flow.post_training_step_update()
    
    # Log
    if step % 100 == 0:
        print(f"Step {step}: loss={loss.item():.4f}")
```

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

### Initialize Unet Model

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Instantiate the Unet model with specified dimensions and configuration. Requires importing the Unet class.

```python
from rectified_flow_pytorch import Unet

model = Unet(
    dim=64,
    channels=3,
    dim_mults=(1, 2, 4, 8),
    depth=4
)
```

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

### Convert Tensors to Float32

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

Resolve 'expected scalar type Double but found Float' errors by ensuring tensor data types match. This example converts input data to float32.

```python
# Convert to float32
data = data.float()
```

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

### PyTorch Lightning Module for Rectified Flow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Integrate Rectified Flow models into PyTorch Lightning for simplified training loops and experiment management. Ensure PyTorch Lightning and Torch are installed.

```python
import pytorch_lightning as pl
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

class RFLitModule(pl.LightningModule):
    def __init__(self):
        super().__init__()
        self.model = Unet(dim=64)
        self.rectified_flow = RectifiedFlow(self.model)
    
    def training_step(self, batch, batch_idx):
        loss = self.rectified_flow(batch)
        self.log('train_loss', loss)
        return loss
    
    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=3e-4)

# Usage
module = RFLitModule()
trainer = pl.Trainer(max_epochs=100, gpus=1)
trainer.fit(module, train_dataloader)
```

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

### load

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Loads the model, optimizer, and EMA state from a specified checkpoint file.

```APIDOC
## load

### Description
Loads the model, optimizer, and EMA state from a specified checkpoint file.

### Method
`load(path: str)`

### Parameters

- **path** (str) - Required - Relative path within checkpoints_folder
```

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

### LsdFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Configure LsdFlow (Laplacian-Scaled) with diagonal probability and options for learned loss weight.

```python
from rectified_flow_pytorch import LsdFlow

lsd_flow = LsdFlow(
    model,
    diag_prob=0.75,
    use_learned_loss_weight=False,
    accept_cond=False
)
```

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

### Conditional Generation with Rectified Flow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

This example demonstrates how to perform conditional generation using Rectified Flow. It involves creating a Unet model that accepts conditioning (`accept_cond=True`) and then providing conditioning embeddings during the sampling process.

```python
import torch
from rectified_flow_pytorch import RectifiedFlow, Unet

# Model with conditioning
model = Unet(
    dim=64,
    accept_cond=True,
    dim_cond=128
)

rectified_flow = RectifiedFlow(model)

# Generate with conditions
batch_size = 16
conditions = torch.randn(batch_size, 128)  # Class embeddings

samples = rectified_flow.sample(
    batch_size=batch_size,
    steps=32,
    data_shape=(3, 256, 256),
    cond=conditions
)
```

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

### NanoFlow Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize NanoFlow for minimal flow matching without ODE solvers. Use it for training loss computation and sample generation.

```python
from rectified_flow_pytorch import NanoFlow

nano_flow = NanoFlow(model, times_cond_kwarg='times')
loss = nano_flow(data)
samples = nano_flow.sample(batch_size=8, steps=16)
```

```python
NanoFlow(
    model: Module,
    times_cond_kwarg: str | None = None,
    data_shape: tuple[int, ...] | None = None,
    normalize_data_fn: Callable = identity,
    unnormalize_data_fn: Callable = identity,
    predict_clean: bool = False,
    max_timesteps: int = 100,
    loss_fn: Callable = F.mse_loss
)
```

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

### Reduce Batch Size for CUDA Memory Errors

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

When encountering CUDA out of memory errors, reduce the batch size to fit within GPU memory. This example shows how to set a smaller batch size for the Trainer.

```python
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    batch_size=8  # Reduced from 32
)
```

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

### Import Standard Pattern

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/module-exports.md

Import the core components for Rectified Flow models.

```python
from rectified_flow_pytorch import Rectified Flow, Unet, Trainer
```

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

### SplitMeanFlow Initialization

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize SplitMeanFlow, a variant of MeanFlow that separates flow-matching and boundary condition objectives.

```python
from rectified_flow_pytorch import SplitMeanFlow

split_mean_flow = SplitMeanFlow(
    model,
    prob_default_flow_obj=0.5
)
```

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

### NanoFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initialize NanoFlow with minimal implementation settings, including data shape and prediction behavior.

```python
from rectified_flow_pytorch import NanoFlow

nano_flow = NanoFlow(
    model,
    times_cond_kwarg='times',
    data_shape=(3, 256, 256),
    normalize_data_fn=identity,
    unnormalize_data_fn=identity,
    predict_clean=False,
    max_timesteps=100
)
```

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

### Basic Rectified Flow Model Training and Sampling

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/00-START-HERE.txt

This snippet demonstrates the fundamental steps for initializing a Rectified Flow model, loading data, training, and then sampling from the trained model. It requires the Unet, RectifiedFlow, ImageDataset, and Trainer classes.

```python
from rectified_flow import RectifiedFlow
from unet import Unet
from dataset import ImageDataset
from trainer import Trainer

# Create model
model = Unet(dim=64, dim_mults=(1, 2, 4, 8))
rf = RectifiedFlow(model)

# Load data
dataset = ImageDataset('./images', image_size=256)

# Train
trainer = Trainer(rf, dataset=dataset, num_train_steps=100_000)
trainer()

# Sample
samples = rf.sample(batch_size=8, steps=32)
```

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

### Use Gradient Accumulation for Larger Effective Batch Size

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/errors.md

If reducing batch size impacts training stability, use gradient accumulation to achieve a larger effective batch size without increasing memory usage. This example sets `grad_accum_every`.

```python
trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    batch_size=8,
    grad_accum_every=4  # Effective batch = 32
)
```

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

### Accelerate Distributed Training Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/types.md

Provides configuration options for Accelerate distributed training, including mixed precision, CPU usage, and gradient accumulation.

```python
{
    'mixed_precision': str = 'no',  # 'no', 'fp16', 'bf16'
    'cpu': bool = False,
    'device_placement': bool = True,
    'gradient_accumulation_steps': int = 1,
    'split_batches': bool = False,
}
```

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

### LsdFlow Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize LsdFlow for Laplacian-scaled diagonal flow, with an option for learned loss weighting. Shows loss computation and sampling.

```python
from rectified_flow_pytorch import LsdFlow

lsd_flow = LsdFlow(
    model,
    diag_prob=0.75,
    use_learned_loss_weight=False
)

loss = lsd_flow(data)
samples = lsd_flow.sample(batch_size=8)
```

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

### Trainer Load Method

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Loads the model, optimizer, and EMA state from a specified checkpoint file. The `path` is relative to the `checkpoints_folder`.

```python
def load(self, path: str) -> None:
    # Load model, optimizer, and EMA state from checkpoint.
    # Parameters:
    # - path (str): Relative path within checkpoints_folder
```

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

### Enable Consistency Flow Matching

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/api-usage.md

Use `use_consistency=True` for faster training with improved quality. Configure ODE solver with `odeint_kwargs`. EMA model can be used for sampling.

```python
from rectified_flow_pytorch import RectifiedFlow, Unet, Trainer, ImageDataset

model = Unet(dim=128)

rectified_flow = RectifiedFlow(
    model,
    use_consistency=True,
    consistency_decay=0.9999,
    odeint_kwargs=dict(atol=1e-5, rtol=1e-5, method='midpoint')
)

dataset = ImageDataset('./images', image_size=256)

trainer = Trainer(
    rectified_flow,
    dataset=dataset,
    num_train_steps=100_000,
    batch_size=32
)

trainer()

# Sample uses EMA model if requested
samples = rectified_flow.sample(batch_size=8, use_ema=True)
```

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

### SelfFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Configure SelfFlow (Teacher-Student) with a FiT model, mask ratio, and alignment layers.

```python
from rectified_flow_pytorch import SelfFlow, FiT

fit = FiT(dim=768, patch_size=16)

self_flow = SelfFlow(
    fit,
    mask_ratio=0.5,
    repr_loss_weight=1.0,
    student_align_layer=-2,
    teacher_align_layer=-1,
    schedule_fn=logit_normal_schedule
)
```

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

### Distributed Training Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Configures the trainer for distributed training using accelerate. Supports mixed precision and device placement.

```python
trainer = Trainer(
    model,
    dataset=dataset,
    accelerate_kwargs=dict(
        mixed_precision='fp16',  # or 'bf16' or 'no'
        cpu=False,
        device_placement=True,
        gradient_accumulation_steps=2
    )
)
```

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

### Configure Built-in Loss Functions for RectifiedFlow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Select and configure built-in loss functions like MSE, Pseudo-Huber, or Pseudo-Huber with LPIPS. Pseudo-Huber offers robustness, while LPIPS enhances perceptual quality.

```python
rectified_flow = RectifiedFlow(model, loss_fn='mse')
```

```python
rectified_flow = RectifiedFlow(
    model,
    loss_fn='pseudo_huber',
    loss_fn_kwargs={'data_dim': 3}
)
```

```python
rectified_flow = RectifiedFlow(
    model,
    loss_fn='pseudo_huber_with_lpips',
    loss_fn_kwargs={
        'data_dim': 3,
        'lpips_kwargs': {}
    }
)
```

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

### MeanFlow Initialization and Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initialize MeanFlow for one-step generative modeling. Supports adaptive loss weighting and logit-normal sampling. Training loss and sampling are demonstrated.

```python
from rectified_flow_pytorch import MeanFlow

mean_flow = MeanFlow(
    model,
    use_adaptive_loss_weight=True,
    use_logit_normal_sampler=True,
    prob_default_flow_obj=0.5
)

loss = mean_flow(data)
samples = mean_flow.sample(batch_size=8)  # One-step by default
```

```python
MeanFlow(
    model: Module,
    data_shape: tuple[int, ...] | None = None,
    normalize_data_fn: Callable = identity,
    unnormalize_data_fn: Callable = identity,
    use_adaptive_loss_weight: bool = True,
    adaptive_loss_weight_p: float = 0.5,
    use_logit_normal_sampler: bool = True,
    logit_normal_mean: float = -0.4,
    logit_normal_std: float = 1.,
    prob_default_flow_obj: float = 0.5,
    add_recon_loss: bool = False,
    recon_loss_weight: float = 1.,
    accept_cond: bool = False,
    noise_std_dev: float = 1.,
    eps: float = 1e-3
)
```

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

### Reflow and ReflowTrainer Usage

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Demonstrates initializing Reflow from an existing RectifiedFlow model and using ReflowTrainer for iterative model refinement.

```python
from rectified_flow_pytorch import Reflow, ReflowTrainer

# Create initial rectified flow
rectified_flow = RectifiedFlow(model)
# Train it...

# Reflow
reflow = Reflow(rectified_flow)
loss = reflow()  # Forward pass generates sample from frozen model, trains on it

# Or use trainer
trainer = ReflowTrainer(
    rectified_flow,
    num_train_steps=100_000
)
trainer()
```

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

### Documentation Structure Overview

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/README.md

This diagram illustrates the hierarchical structure of the documentation, showing the main README file and its related reference, type, configuration, error, API usage, and module export documents.

```markdown
README.md (you are here)
    ├── rectified-flow-reference.md (main API)
    ├── types.md (type definitions)
    ├── configuration.md (configuration options)
    ├── errors.md (troubleshooting)
    ├── api-usage.md (usage examples)
    └── module-exports.md (module structure)
```

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

### ImageDataset Constructor

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Initializes the ImageDataset with a folder path, target image size, and optional parameters for file extensions, horizontal flipping, and image conversion.

```APIDOC
## ImageDataset

PyTorch Dataset for loading images from a folder hierarchy.

### Constructor Signature

```python
ImageDataset(
    folder: str | Path,
    image_size: int,
    exts: list[str] = ['jpg', 'jpeg', 'png', 'tiff'],
    augment_horizontal_flip: bool = False,
    convert_image_to: str | None = None
)
```

### Parameters

#### Parameters

- **folder** (str | Path) - Required - Root folder to search for images
- **image_size** (int) - Required - Target size (images resized and center-cropped)
- **exts** (list[str]) - Optional - Image extensions to search for (glob pattern). Defaults to ['jpg','jpeg','png','tiff'].
- **augment_horizontal_flip** (bool) - Optional - Apply random horizontal flips. Defaults to False.
- **convert_image_to** (str | None) - Optional - Convert all images to this mode (e.g., 'RGB', 'RGBA'). Defaults to None.
```

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

### ODE Integrator Keyword Arguments

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/types.md

Configuration dictionary for `torchdiffeq.odeint` solver. Allows customization of absolute and relative tolerances, and the solver method.

```python
{
    'atol': float = 1e-5,        # Absolute tolerance
    'rtol': float = 1e-5,        # Relative tolerance
    'method': str = 'midpoint'   # ODE solver method
}
```

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

### Configure Probabilistic Output

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Set up Rectified Flow for probabilistic models by enabling `mean_variance_net`. This automatically selects `MeanVarianceNetLoss`. Temperature can be adjusted during sampling for stochasticity.

```python
rectified_flow = RectifiedFlow(
    model,
    mean_variance_net=True,
    loss_fn=MeanVarianceNetLoss()  # Auto-selected if True
)

# During sampling with temperature
samples = rectified_flow.sample(
    batch_size=16,
    temperature=1.5  # Higher = more stochastic
)
```

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

### Configure Unet Basic Dimensions

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Initialize a Unet model with specified dimensions and resolution multipliers. Different configurations are shown for small, medium, and large models.

```python
from rectified_flow_pytorch import Unet

# Small model
model = Unet(
    dim=64,
    dim_mults=(1, 2, 4, 8),
    channels=3
)

# Medium model
model = Unet(
    dim=128,
    dim_mults=(1, 2, 4, 8),
    channels=3
)

# Large model
model = Unet(
    dim=256,
    dim_mults=(1, 2, 4, 8),
    channels=3
)
```

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

### MeanFlow Configuration

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Configure MeanFlow for one-step flow with options for adaptive loss weighting and logit normal sampling.

```python
from rectified_flow_pytorch import MeanFlow

mean_flow = MeanFlow(
    model,
    use_adaptive_loss_weight=True,
    adaptive_loss_weight_p=0.5,
    use_logit_normal_sampler=True,
    logit_normal_mean=-0.4,
    logit_normal_std=1.,
    prob_default_flow_obj=0.5,
    accept_cond=False
)
```

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

### ImageDataset Constructor Signature

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/rectified-flow-reference.md

Defines the parameters available for initializing the ImageDataset, including folder path, image size, extensions, and augmentation options.

```python
ImageDataset(
    folder: str | Path,
    image_size: int,
    exts: list[str] = ['jpg', 'jpeg', 'png', 'tiff'],
    augment_horizontal_flip: bool = False,
    convert_image_to: str | None = None
)
```

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

### Configure ODE Solver Parameters for RectifiedFlow

Source: https://github.com/lucidrains/rectified-flow-pytorch/blob/main/_autodocs/configuration.md

Control ODE solver behavior during sampling using `odeint_kwargs`, specifying tolerances (atol, rtol) and the integration method. Different methods offer trade-offs between speed and accuracy.

```python
rectified_flow = RectifiedFlow(
    model,
    odeint_kwargs=dict(
        atol=1e-5,        # Absolute tolerance
        rtol=1e-5,        # Relative tolerance
        method='midpoint' # ODE solver
    )
)
```