diffusion_model.py

# -*- coding: utf-8 -*-
"""Copy of diffusion_model.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1NyriYRt8prljGy6uaWq4MD0oYu-B8fh1

# A Diffusion Model from Scratch in Pytorch

In this notebook I want to build a very simple (as few code as possible) Diffusion Model for generating car images. I will explain all the theoretical details in the YouTube video. 


**Sources:**
- Github implementation [Denoising Diffusion Pytorch](https://github.com/lucidrains/denoising-diffusion-pytorch)
- Niels Rogge, Kashif Rasul, [Huggingface notebook](https://colab.research.google.com/github/huggingface/notebooks/blob/main/examples/annotated_diffusion.ipynb#scrollTo=3a159023)
- Papers on Diffusion models ([Dhariwal, Nichol, 2021], [Ho et al., 2020] ect.)

## Investigating the dataset

As dataset we use the StandordCars Dataset, which consists of around 8000 images in the train set. Let's see if this is enough to get good results ;-)
"""

# from google.colab import drive
# drive.mount('/content/drive')

# Commented out IPython magic to ensure Python compatibility.
# %cd /content/drive/MyDrive/CV/Shadowremoval

# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/
# !ls
# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours
# !ls
# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A
# !ls
# %cd /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official
# !ls

# !mkdir images
# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_A/images/
# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_A/images/
# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/train/train_C_fixed_ours/images/
# !mv /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official/*.png /content/drive/MyDrive/CV/Shadowremoval/ISTD_Dataset/test/test_C_fixed_official/images/

import torch
import torchvision
import matplotlib.pyplot as plt
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import numpy as np
import torch.nn.functional as F

PATH = {
    'train_shadow' : './ISTD_Dataset/train/train_A', 
    'train_out' : './ISTD_Dataset/train/train_C_fixed_ours', 
    'test_shadow' : './ISTD_Dataset/test', 
    'test_out' : './ISTD_Dataset/test/test_C_fixed_official', 
    }
OUT_PATH = './output/'
IMG_SIZE = 64

data_transform = transforms.Compose([
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.ToTensor(), # Scales data into [0,1] 
        transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1] 
    ])

dataset_train_shadow = datasets.ImageFolder(PATH['train_shadow'], transform=data_transform)
dataset_train_out = datasets.ImageFolder(PATH['train_out'], transform=data_transform)
dataset_test_shadow = datasets.ImageFolder(PATH['test_shadow'], transform=data_transform)
dataset_test_out = datasets.ImageFolder(PATH['test_out'], transform=data_transform)

reverse_transform = transforms.Compose([
        transforms.Lambda(lambda t: (t + 1) / 2),
        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
        transforms.Lambda(lambda t: t * 255.),
        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
        transforms.ToPILImage(),
    ])

INDEX = 0
plt.figure(figsize=(10,10)) 
plt.subplot(2,2,1)
plt.imshow(reverse_transform(dataset_train_shadow[INDEX][0]))
plt.subplot(2,2,2)
plt.imshow(reverse_transform(dataset_train_out[INDEX][0]))
plt.subplot(2,2,3)
plt.imshow(reverse_transform(dataset_test_shadow[INDEX][0]))
plt.subplot(2,2,4)
plt.imshow(reverse_transform(dataset_test_out[INDEX][0]))

print(dataset_test_out[0][0].shape)

"""Later in this notebook we will do some additional modifications to this dataset, for example make the images smaller, convert them to tensors ect.

# Building the Diffusion Model

## Step 1: The forward process = Noise scheduler

We first need to build the inputs for our model, which are more and more noisy images. Instead of doing this sequentially, we can use the closed form provided in the papers to calculate the image for any of the timesteps individually. 

**Key Takeaways**:
- The noise-levels/variances can be pre-computed
- There are different types of variance schedules
- We can sample each timestep image independently (Sums of Gaussians is also Gaussian)
- No model is needed in this forward step
"""

import torch.nn.functional as F

def linear_beta_schedule(timesteps, start=0.0001, end=0.02):
    return torch.linspace(start, end, timesteps)

def get_index_from_list(vals, t, x_shape):
    """ 
    Returns a specific index t of a passed list of values vals
    while considering the batch dimension.
    """
    batch_size = t.shape[0]
    out = vals.gather(-1, t.cpu())
    return out.reshape(batch_size, *((1,) * (len(x_shape) - 1))).to(t.device)

def forward_diffusion_sample(x_0, t, device="cpu"):
    """ 
    Takes an image and a timestep as input and 
    returns the noisy version of it
    """
    noise = torch.randn_like(x_0)
    sqrt_alphas_cumprod_t = get_index_from_list(sqrt_alphas_cumprod, t, x_0.shape)
    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
        sqrt_one_minus_alphas_cumprod, t, x_0.shape
    )
    # mean + variance
    # print(sqrt_alphas_cumprod_t.shape)
    # print(x_0.shape)
    # print(sqrt_one_minus_alphas_cumprod_t.shape)
    # print(noise.shape)
    return sqrt_alphas_cumprod_t.to(device) * x_0.to(device) \
    + sqrt_one_minus_alphas_cumprod_t.to(device) * noise.to(device), noise.to(device)


# Define beta schedule
T = 300
betas = linear_beta_schedule(timesteps=T)

# Pre-calculate different terms for closed form
alphas = 1. - betas
alphas_cumprod = torch.cumprod(alphas, axis=0)
alphas_cumprod_prev = F.pad(alphas_cumprod[:-1], (1, 0), value=1.0)
sqrt_recip_alphas = torch.sqrt(1.0 / alphas)
sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = torch.sqrt(1. - alphas_cumprod)
posterior_variance = betas * (1. - alphas_cumprod_prev) / (1. - alphas_cumprod)

"""Let's test it on our dataset ..."""

# dataset = datasets.ImageFolder('path/to/data', transform=transform)

from torchvision import transforms 
from torch.utils.data import DataLoader
import numpy as np

IMG_SIZE = 64
BATCH_SIZE = 128

# def load_transformed_dataset():
#     data_transforms = [
#         transforms.Resize((IMG_SIZE, IMG_SIZE)),
#         transforms.RandomHorizontalFlip(),
#         transforms.ToTensor(), # Scales data into [0,1] 
#         transforms.Lambda(lambda t: (t * 2) - 1) # Scale between [-1, 1] 
#     ]
#     data_transform = transforms.Compose(data_transforms)

#     train = torchvision.datasets.StanfordCars(root=".", download=True, 
#                                          transform=data_transform)

#     test = torchvision.datasets.StanfordCars(root=".", download=True, 
#                                          transform=data_transform, split='test')
#     return torch.utils.data.ConcatDataset([train, test])
def show_tensor_image(image):
    reverse_transforms = transforms.Compose([
        transforms.Lambda(lambda t: (t + 1) / 2),
        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # CHW to HWC
        transforms.Lambda(lambda t: t * 255.),
        transforms.Lambda(lambda t: t.numpy().astype(np.uint8)),
        transforms.ToPILImage(),
    ])

    # Take first image of batch
    if len(image.shape) == 4:
        image = image[0, :, :, :] 
    plt.imshow(reverse_transforms(image))

# data = load_transformed_dataset()
dataloader_train_shadow = DataLoader(dataset_train_shadow, batch_size=BATCH_SIZE, drop_last=True)
dataloader_train_out = DataLoader(dataset_train_out, batch_size=BATCH_SIZE, drop_last=True)
dataloader_test_shadow = DataLoader(dataset_test_shadow, batch_size=BATCH_SIZE, drop_last=True)
dataloader_test_out = DataLoader(dataset_test_out, batch_size=BATCH_SIZE, drop_last=True)

# Simulate forward diffusion
image = next(iter(dataloader_train_shadow))[0]

plt.figure(figsize=(15,15))
plt.axis('off')
num_images = 10
stepsize = int(T/num_images)

for idx in range(0, T, stepsize):
    t = torch.Tensor([idx]).type(torch.int64)
    plt.subplot(1, num_images+1, (idx//stepsize) + 1)
    image, noise = forward_diffusion_sample(image, t)
    show_tensor_image(image)

"""## Step 2: The backward process = U-Net


"""



"""For a great introduction to UNets, have a look at this post: https://amaarora.github.io/2020/09/13/unet.html.


**Key Takeaways**:
- We use a simple form of a UNet for to predict the noise in the image
- The input is a noisy image, the ouput the noise in the image
- Because the parameters are shared accross time, we need to tell the network in which timestep we are
- The Timestep is encoded by the transformer Sinusoidal Embedding
- We output one single value (mean), because the variance is fixed

"""

from torch import nn
import torch as th
import math


class Block(nn.Module):
    def __init__(self, in_ch, out_ch, time_emb_dim, up=False):
        super().__init__()
        self.time_mlp =  nn.Linear(time_emb_dim, out_ch)
        if up:
            self.conv1 = nn.Conv2d(2*in_ch, out_ch, 3, padding=1)
            self.transform = nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)
        else:
            self.conv1 = nn.Conv2d(in_ch, out_ch, 3, padding=1)
            self.transform = nn.Conv2d(out_ch, out_ch, 4, 2, 1)
        self.conv2 = nn.Conv2d(out_ch, out_ch, 3, padding=1)
        self.bnorm1 = nn.BatchNorm2d(out_ch)
        self.bnorm2 = nn.BatchNorm2d(out_ch)
        self.relu  = nn.ReLU()
        
    def forward(self, x, t, ):
        # First Conv
        h = self.bnorm1(self.relu(self.conv1(x)))
        # Time embedding
        time_emb = self.relu(self.time_mlp(t))
        # Extend last 2 dimensions
        time_emb = time_emb[(..., ) + (None, ) * 2]
        # Add time channel
        h = h + time_emb
        # Second Conv
        h = self.bnorm2(self.relu(self.conv2(h)))
        # Down or Upsample
        return self.transform(h)


class SinusoidalPositionEmbeddings(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, time):
        device = time.device
        half_dim = self.dim // 2
        embeddings = math.log(10000) / (half_dim - 1)
        embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings)
        embeddings = time[:, None] * embeddings[None, :]
        embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1)
        # TODO: Double check the ordering here
        return embeddings


class SimpleUnet(nn.Module):
    """
    A simplified variant of the Unet architecture.
    """
    def __init__(self):
        super().__init__()
        image_channels = 3 * 2
        down_channels = (64, 128, 256, 512, 1024)
        up_channels = (1024, 512, 256, 128, 64)
        out_dim = 1 
        time_emb_dim = 32

        # Time embedding
        self.time_mlp = nn.Sequential(
                SinusoidalPositionEmbeddings(time_emb_dim),
                nn.Linear(time_emb_dim, time_emb_dim),
                nn.ReLU()
            )
        
        # Initial projection
        self.conv0 = nn.Conv2d(image_channels, down_channels[0], 3, padding=1)

        # Downsample
        self.downs = nn.ModuleList([Block(down_channels[i], down_channels[i+1], \
                                    time_emb_dim) \
                    for i in range(len(down_channels)-1)])
        # Upsample
        self.ups = nn.ModuleList([Block(up_channels[i], up_channels[i+1], \
                                        time_emb_dim, up=True) \
                    for i in range(len(up_channels)-1)])

        self.output = nn.Conv2d(up_channels[-1], 3, out_dim)

    def forward(self, x, x1, timestep):
        
        x = th.cat((x.to(device), x1.to(device)), axis = 1).to(device)
        # Embedd time
        t = self.time_mlp(timestep)
        

        # Initial conv
        x = self.conv0(x)
        # Unet
        residual_inputs = []
        for down in self.downs:
            x = down(x, t)
            residual_inputs.append(x)
        for up in self.ups:
            residual_x = residual_inputs.pop()
            # Add residual x as additional channels
            x = torch.cat((x, residual_x), dim=1)           
            x = up(x, t)
        return self.output(x)

model = SimpleUnet()
print("Num params: ", sum(p.numel() for p in model.parameters()))
model

"""**Further improvements that can be implemented:**
- Residual connections
- Different activation functions like SiLU, GWLU, ...
- BatchNormalization 
- GroupNormalization
- Attention
- ...

## Step 3: The loss

**Key Takeaways:**
- After some maths we end up with a very simple loss function
- There are other possible choices like L2 loss ect.
"""

class ChromaticityConsistencyLoss(nn.Module):
    def __init__(self):
        super(ChromaticityConsistencyLoss, self).__init__()

    def forward(self, input_image, output_image):
        input_chromaticity = input_image[0:3] / (input_image[0] + input_image[1] + input_image[2] + 1e-8)
        output_chromaticity = output_image[0:3] / (output_image[0] + output_image[1] + output_image[2] + 1e-8)

        chromaticity_consistency_loss = torch.mean((input_chromaticity - output_chromaticity) ** 2)

        return chromaticity_consistency_loss


# Create instance of ChromaticityConsistencyLoss
chromaticity_consistency_loss = ChromaticityConsistencyLoss()

# Calculate chromaticity consistency loss
# loss = chromaticity_consistency_loss(input_image, output_image)

class StructurePreservationLoss(nn.Module):
    def __init__(self):
        super(StructurePreservationLoss, self).__init__()

    def forward(self, input_image, output_image):
        structure_preservation_loss = torch.mean(torch.abs(input_image - output_image))

        return structure_preservation_loss

def get_loss(model, x_0, x_masked, t):
    x_noisy, noise = forward_diffusion_sample(x_0, t, device)
    noise_pred = model(x_noisy, x_masked, t)
    return F.l1_loss(noise, noise_pred)

"""## Sampling
- Without adding @torch.no_grad() we quickly run out of memory, because pytorch tacks all the previous images for gradient calculation 
- Because we pre-calculated the noise variances for the forward pass, we also have to use them when we sequentially perform the backward process
"""

@torch.no_grad()
def sample_timestep(x, t):
    """
    Calls the model to predict the noise in the image and returns 
    the denoised image. 
    Applies noise to this image, if we are not in the last step yet.
    """
    betas_t = get_index_from_list(betas, t, x.shape)
    sqrt_one_minus_alphas_cumprod_t = get_index_from_list(
        sqrt_one_minus_alphas_cumprod, t, x.shape
    )
    sqrt_recip_alphas_t = get_index_from_list(sqrt_recip_alphas, t, x.shape)
    
    # Call model (current image - noise prediction)
    model_mean = sqrt_recip_alphas_t * (
        x - betas_t * model(x, x, t) / sqrt_one_minus_alphas_cumprod_t
    )
    posterior_variance_t = get_index_from_list(posterior_variance, t, x.shape)
    
    if t == 0:
        return model_mean
    else:
        noise = torch.randn_like(x)
        return model_mean + torch.sqrt(posterior_variance_t) * noise 

@torch.no_grad()
def sample_plot_image(fname):
    # Sample noise
    img_size = IMG_SIZE
    img = torch.randn((1, 3, img_size, img_size), device=device)
    plt.figure(figsize=(15,15))
    plt.axis('off')
    num_images = 10
    stepsize = int(T/num_images)

    for i in range(0,T)[::-1]:
        t = torch.full((1,), i, device=device, dtype=torch.long)
        img = sample_timestep(img, t)
        if i % stepsize == 0:
            plt.subplot(1, num_images, i//stepsize+1)
            show_tensor_image(img.detach().cpu())
    plt.savefig(fname)
    plt.show()

"""## Training"""

from torch.optim import Adam

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
optimizer = Adam(model.parameters(), lr=0.001)
epochs = 10 # Try more!

for epoch in range(epochs):
    for ((step, batch), (step_s, batch_s)) in zip(enumerate(dataloader_train_out), enumerate(dataloader_train_shadow)) :
      optimizer.zero_grad()

      t = torch.randint(0, T, (BATCH_SIZE,), device=device).long()
      loss = get_loss(model, batch[0], batch_s[0], t)

      loss.backward()
      optimizer.step()

      if epoch % 1 == 0 and step == 0:
        print(f"Epoch {epoch} | step {step:03d} Loss: {loss.item()} ")
        fname = OUT_PATH + "Epoch" + str(epoch) + ".jpg"
        sample_plot_image(fname)

"""In Table 2, we show the sample quality effects of reverse process parameterizations and training
objectives (Section 3.2). We find that the baseline option of predicting µ˜ works well only when
trained on the true variational bound instead of unweighted mean squared error, a simplified objective
akin to Eq. (14). We also see that learning reverse process variances (by incorporating a parameterized
diagonal Σθ(xt) into the variational bound) leads to unstable training and poorer sample quality
compared to fixed variances. Predicting , as we proposed, performs approximately as well as
predicting µ˜ when trained on the variational bound with fixed variances, but much better when trained
with our simplified objective.

iffusion models scale down the data with each forward process step (by a √
1 − βt factor)
so that variance does not grow when adding noise, thus providing consistently scaled inputs
to the neural net reverse process. NCSN omits this scaling factor.
"""

torch.save(model.state_dict(), './model.pt')

# Simulate forward diffusion
image, shadow = next(iter(dataloader_train_shadow))[0], next(iter(dataloader_train_shadow))[0]

plt.figure(figsize=(15,15))
plt.axis('off')
num_images = 10
stepsize = int(T/num_images)

for idx in range(0, T, stepsize):
    t = torch.Tensor([idx]).type(torch.int64)
    plt.subplot(1, num_images+1, (idx//stepsize) + 1)
    image, noise = forward_diffusion_sample(image, t)
    show_tensor_image(image)

print(image.shape)

t = torch.randint(0, 5000, (BATCH_SIZE,), device=device).long()
show_tensor_image(model(image, shadow, t).to('cpu').detach())