Training on Celeba-hq about denoising-diffusion-pytorch HOT 5 OPEN

from denoising_diffusion_pytorch import Unet, GaussianDiffusion, Trainer
from torch.utils.data import DataLoader
from torchvision.datasets import VisionDataset
from torchvision.transforms.functional import pil_to_tensor
from PIL import Image

class FFHQDataset(VisionDataset):
    def __init__(self, root: str):
        super().__init__(root)

        self.fpaths = sorted(glob(root + '/**/*.png', recursive=True))
        assert len(self.fpaths) > 0, "File list is empty. Check the root."

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

    def __getitem__(self, index: int):
        fpath = self.fpaths[index]
        img = Image.open(fpath).convert('RGB')
        # normalize to [0, 1] range
        img = pil_to_tensor(img) / 255.
        return img


model = Unet(
    dim = 64,
    dim_mults = (1, 2, 4, 8, 16, 32),
    flash_attn = True
)

diffusion = GaussianDiffusion(
    model,
    image_size = 128,
    timesteps = 1000,           # number of steps
    sampling_timesteps=500
)

dataset = FFHQDataset(root="/mnt/SSD2/nils/ocean_bench_exps/diffusion/data/ffhq/thumbnails128x128")
dataloader = DataLoader(dataset, batch_size = 32, shuffle = True, pin_memory = True, num_workers = 12)

trainer = Trainer(
    diffusion,
    'path/to/your/images',
    train_lr = 8e-5,
    train_num_steps = 50000,         # total training steps
    gradient_accumulate_every = 2,    # gradient accumulation steps
    ema_decay = 0.995,                # exponential moving average decay
    amp = False,
    num_samples=16,
    save_and_sample_every=10000,
    dl = dataloader,
)

trainer.train()