train.py

# Import required libraries
from torchvision import transforms, utils as vutils
from diffusers import UNet2DModel, DDPMScheduler
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from PIL import UnidentifiedImageError
import torch.nn.functional as F
import torch.optim as optim
from tqdm import tqdm
import torch
import os

# Define a safe image folder class to handle errors
class SafeImageFolder(ImageFolder):
    # Override __getitem__ method
    def __getitem__(self, index):
        try:
            # Attempt to get image and label
            return super().__getitem__(index)
        except UnidentifiedImageError:
            # If corrupted, try another random image recursively
            new_index = (index + 1) % len(self)
            return self.__getitem__(new_index)

# Define training function
def train_diffusion():
    # Set hyperparameters
    image_size = 192
    batch_size = 24
    num_epochs = 300
    learning_rate = 1e-4
    save_every = 1
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # Define transforms
    transform = transforms.Compose([
        transforms.Resize(image_size),
        transforms.CenterCrop(image_size),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.ToTensor(),
        transforms.Normalize([0.5] * 3, [0.5] * 3)
    ])

    # Load dataset with safe loader
    dataset = SafeImageFolder("images", transform=transform)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4)

    # Initialize model
    model = UNet2DModel(
        sample_size=image_size,
        in_channels=3,
        out_channels=3,
        layers_per_block=2,
        dropout=0.1,
        block_out_channels=(64, 128, 256, 512),
        down_block_types=("DownBlock2D", "DownBlock2D", "DownBlock2D", "DownBlock2D"),
        up_block_types=("UpBlock2D", "UpBlock2D", "UpBlock2D", "UpBlock2D"),
    ).to(device)

    # Print model architecture
    print(model)

    # Initialize noise scheduler
    noise_scheduler = DDPMScheduler(num_train_timesteps=1000)

    # Initialize optimizer
    optimizer = optim.AdamW(model.parameters(), lr=learning_rate)

    # Initialize cosine annealing learning rate scheduler
    scheduler = optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=num_epochs,
        eta_min=1e-6
    )

    # Initialize gradient scaler
    scaler = torch.amp.GradScaler('cuda')

    # Training loop
    for epoch in range(num_epochs):
        progress_bar = tqdm(dataloader, desc=f"Epoch {epoch + 1}/{num_epochs}")

        for batch in progress_bar:
            # Get real images
            images, _ = batch
            images = images.to(device)

            # Sample noise
            noise = torch.randn_like(images)

            # Sample random timesteps
            timesteps = torch.randint(
                0,
                noise_scheduler.config.num_train_timesteps,
                (images.shape[0],),
                device=device
            ).long()

            # Add noise to images
            noisy_images = noise_scheduler.add_noise(images, noise, timesteps)

            # Enable mixed precision context
            with torch.amp.autocast('cuda'):
                # Predict noise
                noise_pred = model(noisy_images, timesteps).sample

                # Compute loss
                loss = F.mse_loss(noise_pred, noise)

            # Backpropagate with scaled loss
            optimizer.zero_grad()
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()

            # Update progress bar
            progress_bar.set_postfix({
                "loss": f"{loss.item():.4f}",
                "lr": f"{scheduler.get_last_lr()[0]:.6f}"
            })

        # Step scheduler
        scheduler.step()

        # Save progress sample
        if (epoch + 1) % save_every == 0 or epoch == 0:
            save_sample(model, noise_scheduler, epoch + 1, device)

    # Save final model
    model.save_pretrained("unet_final")
    noise_scheduler.save_pretrained("scheduler_final")

# Define sampling function
def save_sample(model, scheduler, epoch, device, num_images=8):
    # Generate noise
    images = torch.randn((num_images, 3, 192, 192), device=device)

    # Set model to evaluation mode
    model.eval()

    # Denoise step-by-step
    with torch.no_grad():
        for t in reversed(range(scheduler.config.num_train_timesteps)):
            # Create a tensor of timesteps
            timesteps = torch.full((num_images,), t, device=device, dtype=torch.long)

            # Predict noise
            with torch.amp.autocast('cuda'):
                noise_pred = model(images, timesteps).sample

            # Step through the scheduler one image at a time
            new_images = []
            for i in range(num_images):
                step_output = scheduler.step(
                    noise_pred[i].unsqueeze(0),
                    timesteps[i].unsqueeze(0),
                    images[i].unsqueeze(0)
                )
                new_images.append(step_output.prev_sample)

            # Stack the new images
            images = torch.cat(new_images, dim=0)

        # Save final denoised images
        vutils.save_image(
            images,
            f"progress/sample_epoch_{epoch:03d}.png",
            nrow=4,
            normalize=True,
            value_range=(-1, 1)
        )

# Define main function
def main():
    # Create progress directory
    os.makedirs("progress", exist_ok=True)

    # Start training
    train_diffusion()

# Run script
if __name__ == "__main__":
    main()