bit-diffusion

The Trainer class provides a high-level interface for training the model on a folder of images. It handles data loading, gradient accumulation, EMA, and periodic sampling/saving of results.

Key parameters for Trainer:

• bit_diffusion: The initialized BitDiffusion instance.
• data_folder: Path to your folder of images.
• results_folder: Directory where training results and samples will be saved.
• num_samples: Number of samples to generate during periodic saving.
• train_batch_size: Training batch size.
• gradient_accumulate_every: Number of steps to accumulate gradients before updating.
• train_lr: Learning rate.
• save_and_sample_every: Frequency (in steps) of saving checkpoints and generating samples.
• train_num_steps: Total number of training steps.
• ema_decay: Exponential moving average decay rate.

from bit_diffusion import Unet, Trainer, BitDiffusion

model = Unet(
    dim = 32,
    channels = 3,
    dim_mults = (1, 2, 4, 8),
).cuda()

bit_diffusion = BitDiffusion(
    model,
    image_size = 128,
    timesteps = 100,
    time_difference = 0.1,       # > 0 helps FID at lower timesteps; set to 0 as timesteps increase
    use_ddim = True              # use ddim
).cuda()

trainer = Trainer(
    bit_diffusion,
    '/path/to/your/data',
    results_folder = './results',
    num_samples = 16,
    train_batch_size = 4,
    gradient_accumulate_every = 4,
    train_lr = 1e-4,
    save_and_sample_every = 1000,
    train_num_steps = 700000,
    ema_decay = 0.995,
)

trainer.train()

You can use the Unet and BitDiffusion classes directly for custom training loops or inference.

Training Note: Input images must be normalized from 0 to 1.

Sampling: After training, use the .sample() method to generate new images from noise.

Key parameters for BitDiffusion:

• model: An instance of Unet.
• image_size: The resolution of the images (e.g., 128).
• timesteps: Number of diffusion timesteps.
• time_difference: A value used during sampling. Research suggests that at lower timesteps, a value greater than 0 helps FID. As timesteps increase, this can be set to 0.
• use_ddim: Boolean to enable DDIM sampling.

import torch
from bit_diffusion import Unet, BitDiffusion

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

bit_diffusion = BitDiffusion(
    model,
    image_size = 128,
    timesteps = 1000
)

# Manual training step
training_images = torch.randn(8, 3, 128, 128) # images are normalized from 0 to 1
loss = bit_diffusion(training_images)
loss.backward()

# After training, sample images
sampled_images = bit_diffusion.sample(batch_size = 4)
sampled_images.shape # (4, 3, 128, 128)