Tutorial 3: Image, Conditional Diffusion and Beyond

Tutorial 3: Image, Conditional Diffusion and Beyond#

Week 2, Day 4: Name of the day

By Neuromatch Academy

Content creators: Binxu Wang

Content reviewers: Shaonan Wang, Dongrui Deng, Dora Zhiyu Yang, Adrita Das

Content editors: Shaonan Wang

Production editors: Spiros Chavlis

Tutorial Objectives#

Understand the idea behind Diffusion generative models: score and reversal of diffusion process.
Learn the score function by denoising data.
Hands-on experience in learning the score to generate certain distributions.

Setup#

Install dependencies#

WARNING: There may be errors and/or warnings reported during the installation. However, they are to be ignored.

Install and import feedback gadget#

# Imports
import random
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import functools

from torch.optim import Adam
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from tqdm.notebook import trange, tqdm
from torch.optim.lr_scheduler import MultiplicativeLR, LambdaLR
from torchvision.utils import make_grid

Figure settings#

Set random seed#

Executing set_seed(seed=seed) you are setting the seed

Set device (GPU or CPU). Execute `set_device()`#

DEVICE = set_device()
SEED = 2021
set_seed(seed=SEED)

WARNING: For this notebook to perform best, if possible, in the menu under `Runtime` -> `Change runtime type.`  select `GPU` 
Random seed 2021 has been set.

Neural Network Architecture#

We just learned the basic principles of diffusion models, with the takeaway that the score function allows us to turn pure noise into some interesting data distribution. Further, we will approximate the score function with a neural network via denoising score matching. But when working with images, we need our neural network to ‘play nice’ with them and to reflect the inductive biases we associate with images.

A reasonable choice is to choose the neural network architecture to be that of a U-Net, which is a CNN-like architecture with:

downscaling/upscaling operations that help the network process features of images at different spatial scales.
skip connection as an information highway.

Since the score function we’re trying to learn a function of time, we also need to devise a way to ensure our neural network properly responds to changes in time. For this purpose, we can use a time embedding.

Video 1: Network architecture#

Submit your feedback#

Coding Exercise 1: Train Diffusion for MNIST#

Finally, let’s implement and train an actual image diffusion model for the MNIST dataset.

By examining the neural network architecture of the score approximator, you will understand the inductive biases we built in.

In the next cell, you will implement the helper functions for the forward process.

marginal_prob_std for \(\sigma_t\) (note, it’s standard deviation, not the variance)
diffusion_coeff for \(g(t)\)

Math Recap for Forward Processes:

We will use the same forward process (variance exploding SDE) as in the last tutorial, which reads:

(116)#\[\begin{equation} d\mathbf x=g(t)d\mathbf w \end{equation}\]

and we let the diffusion coefficient \(g(t)=\lambda^t\), with \(\lambda > 1\).

If so, the marginal distribution of state \(\mathbf x_t\) at time t given an initial state \(\mathbf x_0\) will be a Gaussian \(\mathcal N(\mathbf x_t|\mathbf x_0,\sigma_t^2 I)\). The variance is the integration of the squared diffusion coefficient.

(117)#\[\begin{equation} \sigma_t^2 =\int_0^tg(\tau)^2d\tau=\frac{\lambda^{2t}-1}{2\log\lambda} \end{equation}\]

def marginal_prob_std(t, Lambda, device='cpu'):
  """Compute the standard deviation of $p_{0t}(x(t) | x(0))$.

  Args:
    t: A vector of time steps.
    Lambda: The $\lambda$ in our SDE.

  Returns:
    std : The standard deviation.
  """
  t = t.to(device)
  #################################################
  ## TODO for students: Implement the standard deviation
  raise NotImplementedError("Student exercise: Implement the standard deviation")
  #################################################
  std = ...
  return std


def diffusion_coeff(t, Lambda, device='cpu'):
  """Compute the diffusion coefficient of our SDE.

  Args:
    t: A vector of time steps.
    Lambda: The $\lambda$ in our SDE.

  Returns:
    diff_coeff : The vector of diffusion coefficients.
  """
  #################################################
  ## TODO for students: Implement the diffusion coefficients
  raise NotImplementedError("Student exercise: Implement the diffusion coefficients")
  #################################################
  diff_coeff = ...
  return diff_coeff.to(device)

Click for solution

Submit your feedback#

Network architecture#

Below is code for a simple time embedding and modulation layer. Basically, time \(t\) is multiplexed as sine and cosine basis, then a linear readout creates the time modulation signal.

Time embedding and modulation#

Below is code for a simple U-Net architecture. Apparently, diffusion models can be more or less successful with different architectural details. So this example is mainly for illustrative purposes.

Time-dependent UNet score model#

Show code cell source Hide code cell source

# @title Time-dependent UNet score model

class UNet(nn.Module):
  """A time-dependent score-based model built upon U-Net architecture."""

  def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256], embed_dim=256):
    """Initialize a time-dependent score-based network.

    Args:
      marginal_prob_std: A function that takes time t and gives the standard
        deviation of the perturbation kernel p_{0t}(x(t) | x(0)).
      channels: The number of channels for feature maps of each resolution.
      embed_dim: The dimensionality of Gaussian random feature embeddings.
    """
    super().__init__()
    # Gaussian random feature embedding layer for time
    self.time_embed = nn.Sequential(
          GaussianFourierProjection(embed_dim=embed_dim),
          nn.Linear(embed_dim, embed_dim)
          )
    # Encoding layers where the resolution decreases
    self.conv1 = nn.Conv2d(1, channels[0], 3, stride=1, bias=False)
    self.t_mod1 = Dense(embed_dim, channels[0])
    self.gnorm1 = nn.GroupNorm(4, num_channels=channels[0])

    self.conv2 = nn.Conv2d(channels[0], channels[1], 3, stride=2, bias=False)
    self.t_mod2 = Dense(embed_dim, channels[1])
    self.gnorm2 = nn.GroupNorm(32, num_channels=channels[1])

    self.conv3 = nn.Conv2d(channels[1], channels[2], 3, stride=2, bias=False)
    self.t_mod3 = Dense(embed_dim, channels[2])
    self.gnorm3 = nn.GroupNorm(32, num_channels=channels[2])

    self.conv4 = nn.Conv2d(channels[2], channels[3], 3, stride=2, bias=False)
    self.t_mod4 = Dense(embed_dim, channels[3])
    self.gnorm4 = nn.GroupNorm(32, num_channels=channels[3])


    # Decoding layers where the resolution increases
    self.tconv4 = nn.ConvTranspose2d(channels[3], channels[2], 3, stride=2, bias=False)
    self.t_mod5 = Dense(embed_dim, channels[2])
    self.tgnorm4 = nn.GroupNorm(32, num_channels=channels[2])
    self.tconv3 = nn.ConvTranspose2d(channels[2] + channels[2], channels[1], 3, stride=2, bias=False, output_padding=1)
    self.t_mod6 = Dense(embed_dim, channels[1])
    self.tgnorm3 = nn.GroupNorm(32, num_channels=channels[1])
    self.tconv2 = nn.ConvTranspose2d(channels[1] + channels[1], channels[0], 3, stride=2, bias=False, output_padding=1)
    self.t_mod7 = Dense(embed_dim, channels[0])
    self.tgnorm2 = nn.GroupNorm(32, num_channels=channels[0])
    self.tconv1 = nn.ConvTranspose2d(channels[0] + channels[0], 1, 3, stride=1)

    # The swish activation function
    self.act = lambda x: x * torch.sigmoid(x)
    # A restricted version of the `marginal_prob_std` function, after specifying a Lambda.
    self.marginal_prob_std = marginal_prob_std

  def forward(self, x, t, y=None):
    # Obtain the Gaussian random feature embedding for t
    embed = self.act(self.time_embed(t))
    # Encoding path, downsampling
    ## Incorporate information from t
    h1 = self.conv1(x)  + self.t_mod1(embed)
    ## Group normalization  and  apply activation function
    h1 = self.act(self.gnorm1(h1))
    #  2nd conv
    h2 = self.conv2(h1) + self.t_mod2(embed)
    h2 = self.act(self.gnorm2(h2))
    # 3rd conv
    h3 = self.conv3(h2) + self.t_mod3(embed)
    h3 = self.act(self.gnorm3(h3))
    # 4th conv
    h4 = self.conv4(h3) + self.t_mod4(embed)
    h4 = self.act(self.gnorm4(h4))

    # Decoding path up sampling
    h = self.tconv4(h4) + self.t_mod5(embed)
    ## Skip connection from the encoding path
    h = self.act(self.tgnorm4(h))
    h = self.tconv3(torch.cat([h, h3], dim=1)) + self.t_mod6(embed)
    h = self.act(self.tgnorm3(h))
    h = self.tconv2(torch.cat([h, h2], dim=1)) + self.t_mod7(embed)
    h = self.act(self.tgnorm2(h))
    h = self.tconv1(torch.cat([h, h1], dim=1))

    # Normalize output
    h = h / self.marginal_prob_std(t)[:, None, None, None]
    return h

Think! 1: U-Net Architecture#

Looking at the U-Net architecture, can you find the module(s) corresponding to the following operations?

Downsampling the spatial features?
Upsampling the spatial features?
The skip connection from the down branch to the up branch, how is it implemented?
How is time modulation implemented?
Why is the output divided by self.marginal_prob_std(t) before output? How might this help or harm the score learning?

Take 2 minutes to think in silence, then discuss as a group (~10 minutes).

Click for solution

Submit your feedback#

Coding Exercise 2: Defining the loss function#

In the next cell, you will implement the denoising score matching (DSM) objective as we used in the last tutorial.

(118)#\[\begin{equation} \mathcal L=\int_\epsilon^1dt \mathbb E_{x\sim p_0(x)}\mathbb E_{z\sim \mathcal N(0,I)}\|\sigma_t s_\theta(x+\sigma_t z, t)+z\|^2 \end{equation}\]

where the time weighting is chosen as \(\gamma_t=\sigma_t^2\), which emphasizes the high noise period (\(t\sim 1\)) more than the low noise period (\(t\sim 0\)).

Tips:

The major difference from the last tutorial is that the score \(s\), noise \(z\), and states \(x\) are all batch image-shaped tensor, so remember to broadcast the \(\sigma_t\) properly. e.g. this std[:, None, None, None] will be helpful.
eps is set at a small number to stop the model from learning the score function of a very small noise scale, which is highly irregular.

def loss_fn(model, x, marginal_prob_std, eps=1e-3, device='cpu'):
  """The loss function for training score-based generative models.

  Args:
    model: A PyTorch model instance that represents a
      time-dependent score-based model.
      Note, it takes two inputs in its forward function model(x, t)
      $s_\theta(x,t)$ in the equation
    x: A mini-batch of training data.
    marginal_prob_std: A function that gives the standard deviation of
      the perturbation kernel, takes `t` as input.
      $\sigma_t$ in the equation.
    eps: A tolerance value for numerical stability.
  """
  # Sample time uniformly in eps, 1
  random_t = torch.rand(x.shape[0], device=device) * (1. - eps) + eps
  # Find the noise std at the time `t`
  std = marginal_prob_std(random_t).to(device)
  #################################################
  ## TODO for students: Implement the denoising score matching eq.
  raise NotImplementedError("Student exercise: Implement the denoising score matching eq. ")
  #################################################
  # get normally distributed noise N(0, I)
  z = ...
  # compute the perturbed x = x + z * \sigma_t
  perturbed_x = ...
  # predict score with the model at (perturbed x, t)
  score = ...
  # compute distance between the score and noise \| score * sigma_t + z \|_2^2
  loss = ...
  ##############
  return loss

Click for solution

A correctly implemented loss function shall pass the test below.

For a dataset with a single 0 datapoint, we have the analytical score \(\mathbf s(\mathbf x,t)=-\mathbf x/\sigma_t^2\). We test that, for this case, the analytical has zero loss.

Test loss function#

Submit your feedback#

Train and Test the Diffusion Model#

Note: We have reduced the n_epochs to 12, but feel free to increase and use a larger value. The original value was set to 100, but if the training takes too long, n_epochs=50 with batch_size=1024 also suffice. An average loss of around ~30 can generate acceptable digits.

Training the model#

Define the Sampler#

Sampling#

---------------------------------------------------------------------------
FileNotFoundError                         Traceback (most recent call last)
Cell In[21], line 30
     28 marginal_prob_std_fn = lambda t: marginal_prob_std(t, Lambda=Lambda, device=DEVICE)
     29 uncond_score_model = UNet(marginal_prob_std=marginal_prob_std_fn)
---> 30 uncond_score_model.load_state_dict(torch.load("ckpt.pth"))
     31 uncond_score_model.to(DEVICE)
     32 save_samples_uncond(uncond_score_model, suffix="", device=DEVICE)

File /opt/hostedtoolcache/Python/3.9.21/x64/lib/python3.9/site-packages/torch/serialization.py:998, in load(f, map_location, pickle_module, weights_only, mmap, **pickle_load_args)
    995 if 'encoding' not in pickle_load_args.keys():
    996     pickle_load_args['encoding'] = 'utf-8'
--> 998 with _open_file_like(f, 'rb') as opened_file:
    999     if _is_zipfile(opened_file):
   1000         # The zipfile reader is going to advance the current file position.
   1001         # If we want to actually tail call to torch.jit.load, we need to
   1002         # reset back to the original position.
   1003         orig_position = opened_file.tell()

File /opt/hostedtoolcache/Python/3.9.21/x64/lib/python3.9/site-packages/torch/serialization.py:445, in _open_file_like(name_or_buffer, mode)
    443 def _open_file_like(name_or_buffer, mode):
    444     if _is_path(name_or_buffer):
--> 445         return _open_file(name_or_buffer, mode)
    446     else:
    447         if 'w' in mode:

File /opt/hostedtoolcache/Python/3.9.21/x64/lib/python3.9/site-packages/torch/serialization.py:426, in _open_file.__init__(self, name, mode)
    425 def __init__(self, name, mode):
--> 426     super().__init__(open(name, mode))

FileNotFoundError: [Errno 2] No such file or directory: 'ckpt.pth'

Nice job! You have just finished the training of a Diffusion model. As you see, the result is not ideal, and many factors affect this. To name a few:

Better network architecture: residual connections, attention mechanism, better upsampling mechanism
Better objective: better weighting function \(\gamma_t\)
Better optimization procedure: using learning rate decay
Better sampling algorithm: Euler integration is known to have larger errors, so it’s advisable to use a more advanced SDE or ODE solver

Section 2: Conditional Diffusion Model#

Another way to greatly improve the result is adding a conditional signal – for example, tell the score network which digit you want. This makes the score modeling much more effortless and adds controllability to the user. The popular Stable Diffusion model is one of this kind, which uses natural language text as the conditional signal for images.

Video 2: Conditional Diffusion Model#

Submit your feedback#

In formulation, the conditional diffusion is highly similar to the unconditional diffusion.

If you are curious about how to build and train a conditional diffusion model, you are welcome to look at the Bonus exercise at the end.

Video 3: Advanced Techinque - Stable Diffusion#

Submit your feedback#

Interactive Demo 2: Stable Diffusion#

In this demo, we will play with one of the most potent open-source diffusion models, Stable Diffusion 2.1, and try to connect with what we have learned.

Download the Stable Diffusion models#

Now you can let loose your imagination and create artworks from text!

Example prompts:

prompt = "A lovely cat running on the dessert in Van Gogh style, trending art."
prompt = "A ballerina dancing under the starry night in Monet style, trending art."

prompt = "A lovely cat running on the dessert in Van Gogh style, trending art." # @param {'type':'string'}
my_seed = 2023  # @param {'type':'integer'}
execute = False  # @param {'type':'boolean'}

if execute:
  image = pipe(prompt, num_inference_steps=50,
              generator=torch.Generator("cuda").manual_seed(my_seed)).images[0]
  image

Submit your feedback#

Think! 2: Architecture of Stable Diffusion Model#

Can you see the similarity between the U-Net in Stable Diffusion and the baby UNet we defined up there? To inspect the architecture, you can use the recursive_print(pipe.unet,deepest=2) function with a different deepest.

The text is encoded through the CLIP model, and you can also look at its structure below recursive_print(pipe.text_encoder,deepest=4), which is a large transformer!

Take 2 minutes to think and play with the code, then discuss as a group (~10 minutes).

Helper function to inspect network#

recursive_print(pipe.unet,deepest=2)

[UNet2DConditionModel]
(conv_in): Conv2d(4, 320, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(time_proj): Timesteps()
(time_embedding): TimestepEmbedding
  (linear_1): Linear(in_features=320, out_features=1280, bias=True)
  (act): SiLU()
  (linear_2): Linear(in_features=1280, out_features=1280, bias=True)
(down_blocks): ModuleList len=4
  (0): CrossAttnDownBlock2D
  (1): CrossAttnDownBlock2D
  (2): CrossAttnDownBlock2D
  (3): DownBlock2D
(up_blocks): ModuleList len=4
  (0): UpBlock2D
  (1): CrossAttnUpBlock2D
  (2): CrossAttnUpBlock2D
  (3): CrossAttnUpBlock2D
(mid_block): UNetMidBlock2DCrossAttn
  (attentions): ModuleList len=1
  (resnets): ModuleList len=2
(conv_norm_out): GroupNorm(32, 320, eps=1e-05, affine=True)
(conv_act): SiLU()
(conv_out): Conv2d(320, 4, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))

recursive_print(pipe.text_encoder,deepest=4)

[CLIPTextModel]
(text_model): CLIPTextTransformer
  (embeddings): CLIPTextEmbeddings
    (token_embedding): Embedding(49408, 1024)
    (position_embedding): Embedding(77, 1024)
  (encoder): CLIPEncoder
    (layers): ModuleList len=23
      (0): CLIPEncoderLayer
      (1): CLIPEncoderLayer
      (2): CLIPEncoderLayer
      (3): CLIPEncoderLayer
      (4): CLIPEncoderLayer
      (5): CLIPEncoderLayer
      (6): CLIPEncoderLayer
      (7): CLIPEncoderLayer
      (8): CLIPEncoderLayer
      (9): CLIPEncoderLayer
      (10): CLIPEncoderLayer
      (11): CLIPEncoderLayer
      (12): CLIPEncoderLayer
      (13): CLIPEncoderLayer
      (14): CLIPEncoderLayer
      (15): CLIPEncoderLayer
      (16): CLIPEncoderLayer
      (17): CLIPEncoderLayer
      (18): CLIPEncoderLayer
      (19): CLIPEncoderLayer
      (20): CLIPEncoderLayer
      (21): CLIPEncoderLayer
      (22): CLIPEncoderLayer
  (final_layer_norm): LayerNorm((1024,), eps=1e-05, elementwise_affine=True)

Click for solution

Submit your feedback#

Section 3: Ethical Considerations#

Video 4: Ethical Consideration#

Submit your feedback#

Think! 3: Copyright of imagery generated from diffusion generated models#

Suppose you prompt a pretrained diffusion model with the name of the artist and obtain beautiful imagery similar to that artist’s style. Who has the copyright of the generated image? The producing company of the diffusion model, the original artist, you, the prompter, the random seed and the weights, or the GPU that runs the inference?

Who do you think deserves the credit and why?

What if you apply enough post-processing steps to the generated images, e.g., finetune the prompt and seed, or edit the image?

Take 2 minutes to think in silence, then discuss as a group (~10 minutes).

Click for solution

Submit your feedback#

Summary#

Today, we learned about

One major application for diffusion modeling, i.e., Modeling natural images.
Inductive biases suitable for image modeling: U-Net architecture and time modulation mechanism.
Introduction to conditional diffusion models, with a demo on Stable Diffusion.
Ethical considerations related to diffusion models, including copyright, misinformation, and fairness.

Daily survey#

Don’t forget to complete your reflections and content check in the daily survey! Please be patient after logging in as there is a small delay before you will be redirected to the survey.

Bonus: Train Conditional Diffusion for MNIST#

In this part, we’d like to train an MNIST generative model conditioned on the digit.

Here we will use a basic form of conditional modulation, i.e., digit embedding, to linearly control the relative gain of the features. After learning about the attention mechanism, you could think about better ways conditional modulation, e.g., using cross-attention to modulate the score model.

UNet score model with conditional modulation#

class UNet_Conditional(nn.Module):
  """A time-dependent score-based model built upon U-Net architecture."""

  def __init__(self, marginal_prob_std, channels=[32, 64, 128, 256], embed_dim=256,
               text_dim=256, nClass=10):
    """Initialize a time-dependent score-based network.

    Args:
      marginal_prob_std: A function that takes time t and gives the standard
        deviation of the perturbation kernel p_{0t}(x(t) | x(0)).
      channels: The number of channels for feature maps of each resolution.
      embed_dim: The dimensionality of Gaussian random feature embeddings of time.
      text_dim:  the embedding dimension of text / digits.
      nClass:    number of classes you want to model.
    """
    super().__init__()
    # random embedding for classes
    self.cond_embed = nn.Embedding(nClass, text_dim)
    # Gaussian random feature embedding layer for time
    self.time_embed = nn.Sequential(
        GaussianFourierProjection(embed_dim=embed_dim),
        nn.Linear(embed_dim, embed_dim)
        )
    # Encoding layers where the resolution decreases
    self.conv1 = nn.Conv2d(1, channels[0], 3, stride=1, bias=False)
    self.t_mod1 = Dense(embed_dim, channels[0])
    self.gnorm1 = nn.GroupNorm(4, num_channels=channels[0])

    self.conv2 = nn.Conv2d(channels[0], channels[1], 3, stride=2, bias=False)
    self.t_mod2 = Dense(embed_dim, channels[1])
    self.gnorm2 = nn.GroupNorm(32, num_channels=channels[1])
    self.y_mod2 = Dense(embed_dim, channels[1])

    self.conv3 = nn.Conv2d(channels[1], channels[2], 3, stride=2, bias=False)
    self.t_mod3 = Dense(embed_dim, channels[2])
    self.gnorm3 = nn.GroupNorm(32, num_channels=channels[2])
    self.y_mod3 = Dense(embed_dim, channels[2])

    self.conv4 = nn.Conv2d(channels[2], channels[3], 3, stride=2, bias=False)
    self.t_mod4 = Dense(embed_dim, channels[3])
    self.gnorm4 = nn.GroupNorm(32, num_channels=channels[3])
    self.y_mod4 = Dense(embed_dim, channels[3])

    # Decoding layers where the resolution increases
    self.tconv4 = nn.ConvTranspose2d(channels[3], channels[2], 3, stride=2, bias=False)
    self.t_mod5 = Dense(embed_dim, channels[2])
    self.y_mod5 = Dense(embed_dim, channels[2])
    self.tgnorm4 = nn.GroupNorm(32, num_channels=channels[2])

    self.tconv3 = nn.ConvTranspose2d(channels[2], channels[1], 3, stride=2, bias=False, output_padding=1)     #  + channels[2]
    self.t_mod6 = Dense(embed_dim, channels[1])
    self.y_mod6 = Dense(embed_dim, channels[1])
    self.tgnorm3 = nn.GroupNorm(32, num_channels=channels[1])

    self.tconv2 = nn.ConvTranspose2d(channels[1], channels[0], 3, stride=2, bias=False, output_padding=1)     #  + channels[1]
    self.t_mod7 = Dense(embed_dim, channels[0])
    self.y_mod7 = Dense(embed_dim, channels[0])
    self.tgnorm2 = nn.GroupNorm(32, num_channels=channels[0])
    self.tconv1 = nn.ConvTranspose2d(channels[0], 1, 3, stride=1)

    # The swish activation function
    self.act = nn.SiLU() # lambda x: x * torch.sigmoid(x)
    self.marginal_prob_std = marginal_prob_std
    for module in [self.y_mod2,self.y_mod3,self.y_mod4,
                   self.y_mod5,self.y_mod6,self.y_mod7]:
        nn.init.normal_(module.dense.weight, mean=0, std=0.0001)
        nn.init.constant_(module.dense.bias, 1.0)

  def forward(self, x, t, y=None):
    # Obtain the Gaussian random feature embedding for t
    embed = self.act(self.time_embed(t))
    y_embed = self.cond_embed(y)
    # Encoding path
    h1 = self.conv1(x) + self.t_mod1(embed)
    ## Incorporate information from t
    ## Group normalization
    h1 = self.act(self.gnorm1(h1))
    h2 = self.conv2(h1) + self.t_mod2(embed)
    h2 = h2 * self.y_mod2(y_embed)
    h2 = self.act(self.gnorm2(h2))
    h3 = self.conv3(h2) + self.t_mod3(embed)
    h3 = h3 * self.y_mod3(y_embed)
    h3 = self.act(self.gnorm3(h3))
    h4 = self.conv4(h3) + self.t_mod4(embed)
    h4 = h4 * self.y_mod4(y_embed)
    h4 = self.act(self.gnorm4(h4))

    # Decoding path
    h = self.tconv4(h4) + self.t_mod5(embed)
    h = h * self.y_mod5(y_embed)
    ## Skip connection from the encoding path
    h = self.act(self.tgnorm4(h))
    h = self.tconv3(h + h3) + self.t_mod6(embed)
    h = h * self.y_mod6(y_embed)
    h = self.act(self.tgnorm3(h))
    h = self.tconv2(h + h2) + self.t_mod7(embed)
    h = h * self.y_mod7(y_embed)
    h = self.act(self.tgnorm2(h))
    h = self.tconv1(h + h1)

    # Normalize output
    h = h / self.marginal_prob_std(t)[:, None, None, None]
    return h

Loss for conditional diffusion#

def loss_fn_cond(model, x, y, marginal_prob_std, eps=1e-3):
  """The loss function for training score-based generative models.

  Args:
    model: A PyTorch model instance that represents a
      time-dependent score-based model.
    x: A mini-batch of training data.
    marginal_prob_std: A function that gives the standard deviation of
      the perturbation kernel.
    eps: A tolerance value for numerical stability.
  """
  random_t = torch.rand(x.shape[0], device=x.device) * (1. - eps) + eps
  z = torch.randn_like(x)
  std = marginal_prob_std(random_t)
  perturbed_x = x + z * std[:, None, None, None]
  score = model(perturbed_x, random_t, y=y)
  loss = torch.mean(torch.sum((score * std[:, None, None, None] + z)**2,
                              dim=(1, 2, 3)))
  return loss

Training conditional diffusion model#

Note: The original value for n_epochs was 100.

Sample Conditional Diffusion#

torch.cuda.empty_cache()

Tutorial 3: Image, Conditional Diffusion and Beyond

Contents

Tutorial 3: Image, Conditional Diffusion and Beyond#

Tutorial Objectives#

Setup#

Install dependencies#

Install and import feedback gadget#

Figure settings#

Set random seed#

Set device (GPU or CPU). Execute set_device()#

Neural Network Architecture#

Video 1: Network architecture#

Submit your feedback#

Coding Exercise 1: Train Diffusion for MNIST#

Submit your feedback#

Network architecture#

Time embedding and modulation#

Time-dependent UNet score model#

Think! 1: U-Net Architecture#

Submit your feedback#

Coding Exercise 2: Defining the loss function#

Test loss function#

Submit your feedback#

Train and Test the Diffusion Model#

Training the model#

Define the Sampler#

Sampling#

Section 2: Conditional Diffusion Model#

Video 2: Conditional Diffusion Model#

Submit your feedback#

Video 3: Advanced Techinque - Stable Diffusion#

Submit your feedback#

Interactive Demo 2: Stable Diffusion#

Download the Stable Diffusion models#

Submit your feedback#

Think! 2: Architecture of Stable Diffusion Model#

Helper function to inspect network#

Submit your feedback#

Section 3: Ethical Considerations#

Video 4: Ethical Consideration#

Submit your feedback#

Think! 3: Copyright of imagery generated from diffusion generated models#

Submit your feedback#

Summary#

Daily survey#

Bonus: Train Conditional Diffusion for MNIST#

UNet score model with conditional modulation#

Loss for conditional diffusion#

Training conditional diffusion model#

Sample Conditional Diffusion#

Set device (GPU or CPU). Execute `set_device()`#