denoising_diffusion.py

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer

if __name__ == '__main__':
    model = Unet(
        dim=64,
        dim_mults=(1, 2, 4, 8)
    )

    diffusion = GaussianDiffusion(
        model,
        image_size=256,
        timesteps=1000,
        loss_type='l1'
    )

    trainer = Trainer(
        diffusion,
        'data/AID_mountains',
        train_batch_size=10,
        train_lr=8e-5,
        train_num_steps=10000,
        gradient_accumulate_every=2,
        ema_decay=0.995,
        amp=False,
        calculate_fid=False,
        convert_image_to='RGB',
        save_and_sample_every=1000,
        num_samples=36
    )

    trainer.train()