"The objective of this tutorial is to discover Diffusion Models that are probabilistic generative models which learn to generate data by iteratively refining random noise through a reverse diffusion process. Given a sample of data, noise is progressively added in small steps until it becomes pure noise. Then, a neural network is trained to reverse this process and generate realistic data from noise. \n",
"\n",
"Diffusion models have gained popularity due to their ability to generate high-quality, diverse, and detailed content, surpassing GANs in the quality of the generated images. \n",
"\n",
"In this assignment we will focus on DDPMs, which were introduced in this [paper](https://arxiv.org/abs/2006.11239) and laid the foundation for generative diffusion models.\n",
"\n",
"The notebook contains code cells with the **\"# TO DO\"** comments. Your goal is to complete these cells and run the proposed experiments. \n",
"\n",
"As the computation is heavy, particularly during training, we encourage you to use a GPU. If your laptob is not equiped, you may use one of these remote jupyter servers, where you can select the execution on GPU :\n",
"This server is accessible within the campus network. If outside, you need to use a VPN. Before executing the notebook, select the kernel \"Python PyTorch\" to run it on GPU and have access to PyTorch module.\n",
"Before executing the notebook, select the execution on GPU : \"Exécution\" Menu -> \"Modifier le type d'exécution\" and select \"T4 GPU\". "
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "M-WNKvhOP1ED"
},
"source": [
"# Part1: Diffusion"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# First import useful libraries\n",
"import torch\n",
"import numpy as np\n",
"import matplotlib.pyplot as plt\n",
"import torchvision.transforms as transforms"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# device = torch.device(\"mps\")\n",
"device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "y_r8nMTGQI9a"
},
"source": [
"As in the previous session, we will use the MNIST dataset."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# TO DO: your code here to load the MNIST dataset. The size of the images should be set to 64x64.\n",
"mnist_dataset = \n",
"mnist_dataloader ="
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "jiHCy4_UUBFb"
},
"source": [
"Auxiliary functions for plotting images"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"def reverse_transform(image):\n",
" image = image.numpy().transpose((1, 2, 0))\n",
" image = np.clip(image, 0, 1)\n",
" image = (image * 255).astype(np.uint8)\n",
"\n",
" return image\n",
"\n",
"def plot1xNArray(images, labels):\n",
" f, axarr = plt.subplots(1, len(images))\n",
" \n",
" for image, ax, label in zip(images, axarr, labels):\n",
" ax.imshow(image, cmap='gray')\n",
" ax.axis('off')\n",
" ax.set_title(label)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "7SjXNoT7BUey"
},
"source": [
"In order to train the model with the diffusion process, we will use a noise scheduler, which will be in charge of the forward diffusion process. The scheduler takes an image, a sample of random noise and a timestep, and return a noisy image for the corresponding timestep. Noise is progressively added to the image at each timestep, therefore a noisy image at timestep 0 will have barely any noise while a noisy image at the maximum timestep will be basically just noise. \n",
"\n",
"Let's create a noise scheduler with 1000 max timesteps and visualize some noise images. \n",
"\n",
"We will use the diffusers library from Hugging Face, which provides several tools for training and using diffusion models. \n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "uOKvYDyu0w8N"
},
"outputs": [],
"source": [
"from diffusers import DDPMScheduler\n",
"\n",
"# TO DO: Create the scheduler\n",
"noise_scheduler = \n",
"\n",
"image, _ = mnist_dataset[0]\n",
"\n",
"# TO DO: Create a noise tensor sampled from a normal distribution with the same shape as the image\n",
"For the reverse diffusion process we will use a neural network. Given a noisy image and the corresponding timestep, the goal of the neural network is to predict the noise, which allows for the denoising. \n",
"\n",
"For the model, we will have a similar architecture as we used for the cGAN generator, a 2D UNet with a few modifications. The main difference will be that we have to indicate to the model which timestep is currently being denoised. For that purpose, a timestep embedding is added, therefore the model has 2 inputs, the noisy image and the corresponding timestep. \n",
" \n",
"In this exercise, we will use an UNet implementation from the diffusers library, which already has the timestep embedding included."
"If the training takes too long, you can download an already trained model from [this link](https://partage.liris.cnrs.fr/index.php/s/AP2t6b3w8SM4Bp5) and use it for inference. "
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {},
"colab_type": "code",
"id": "1hmcejTWJSYY"
},
"outputs": [],
"source": [
"# TO DO: Add the path to the model checkpoint for loading the model\n",
"During training, for each data sample, we take a random timestep and the corresponding noisy image to give it as input to our model. With sufficient training, the model should learn how to predict the noise in a noisy image for all possible timesteps. \n",
"\n",
"During inference, to generate an image, we will start from pure noise and step by step predict the noise to go from one noisy image to the next, progressively denoising the image until we reach the timestep 0, in which we should have an image without any noise."
" plot1xNArray([img[i] for img in images], labels)"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "PhPkU7BDYooV"
},
"source": [
"The diffusers library also provides *Pipeline* classes, which are wrappers around the model that abstracts the inference loop implemented above. \n",
" \n",
"We can create a pipeline, giving it the trained model and the noise scheduler, and use it to generate images. In this case, we will only have access to the final image, generated on the last timestep, but not the intermediary images from the denoising process. "
"In this exercise we achieved decent results for very simple datasets. But we are quite far from those beautiful AI generated images we can find online. That is for 2 main reasons:\n",
"- Model size: due to the computation and time constrains, we can't really train very large models \n",
"- Dataset size: due to the same constrains, we can't use very complex and large datasets, which requires larger models and longer training times.\n",
"\n",
"Fortunately, even though we can't train those large models with the available hardware and time, we can at least use them for inference ! \n",
"\n",
"The goal of this part is to learn how to retrieve and use a pre-trained diffusion models and also to get creative to come up with some nice prompts to generate outstanding images. \n",
"\n",
"We are going to use Stable Diffusion 3.5, which is a state of the art open-source text conditioned model. It takes a prompt in natural language and use it to guide the diffusion process. This type of models are trained with image-text pairs, but can generalize beyond the pairs seen during training, being able to mix several different concepts into a single image. \n",
"\n",
"In order to save memory, we will use quantization, which consists into converting the model weights types from float16 into float4. That simply means that each model weight will be stored using only 4 bits instead of 16 bits. That allows us to run the model in GPUs with less VRAM and have faster inference, with a small drop in the quality of the results."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For this part of the assignment restart the notebook kernel to be sure your GPU memory is empty. The memory usage can be verified using the command `nvidia-smi` in a terminal. If your GPU has 2GB of VRAM or less, the model will probably not fit into memory even with quantization. In that case, use Google Colab for this part or use the smaller model indicated below. If you are not happy with the results and have plenty of VRAM available, feel free to increase the quantization to 8 bits or even load the model without quantization."
"# TO DO: test different prompts and visualize the generated images\n",
"# Once you are happy with the results, you can save 3 different images as png file with the correspondent prompts in a text file\n",
"prompt = \n",
"\n",
"image = pipeline(\n",
" prompt=prompt,\n",
" num_inference_steps=40,\n",
" guidance_scale=4.5,\n",
" max_sequence_length=512\n",
").images[0]\n",
"image.save(\"generated_image.png\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"colab_type": "text",
"id": "Okb3LU76BUgG"
},
"source": [
"If even with the quantization you still run out of GPU memory and you can't use Google Colab, you can use the following code instead, which uses a much smaller model (the results won't be as near as impressive, but it should be able to run even on a CPU, if you have a little patience)"
The objective of this tutorial is to discover Diffusion Models that are probabilistic generative models which learn to generate data by iteratively refining random noise through a reverse diffusion process. Given a sample of data, noise is progressively added in small steps until it becomes pure noise. Then, a neural network is trained to reverse this process and generate realistic data from noise.
Diffusion models have gained popularity due to their ability to generate high-quality, diverse, and detailed content, surpassing GANs in the quality of the generated images.
In this assignment we will focus on DDPMs, which were introduced in this [paper](https://arxiv.org/abs/2006.11239) and laid the foundation for generative diffusion models.
The notebook contains code cells with the **"# TO DO"** comments. Your goal is to complete these cells and run the proposed experiments.
As the computation is heavy, particularly during training, we encourage you to use a GPU. If your laptob is not equiped, you may use one of these remote jupyter servers, where you can select the execution on GPU :
This server is accessible within the campus network. If outside, you need to use a VPN. Before executing the notebook, select the kernel "Python PyTorch" to run it on GPU and have access to PyTorch module.
As in the previous session, we will use the MNIST dataset.
%% Cell type:code id: tags:
``` python
# TO DO: your code here to load the MNIST dataset. The size of the images should be set to 64x64.
mnist_dataset=
mnist_dataloader=
```
%% Cell type:markdown id: tags:
Auxiliary functions for plotting images
%% Cell type:code id: tags:
``` python
defreverse_transform(image):
image=image.numpy().transpose((1,2,0))
image=np.clip(image,0,1)
image=(image*255).astype(np.uint8)
returnimage
defplot1xNArray(images,labels):
f,axarr=plt.subplots(1,len(images))
forimage,ax,labelinzip(images,axarr,labels):
ax.imshow(image,cmap='gray')
ax.axis('off')
ax.set_title(label)
```
%% Cell type:markdown id: tags:
In order to train the model with the diffusion process, we will use a noise scheduler, which will be in charge of the forward diffusion process. The scheduler takes an image, a sample of random noise and a timestep, and return a noisy image for the corresponding timestep. Noise is progressively added to the image at each timestep, therefore a noisy image at timestep 0 will have barely any noise while a noisy image at the maximum timestep will be basically just noise.
Let's create a noise scheduler with 1000 max timesteps and visualize some noise images.
We will use the diffusers library from Hugging Face, which provides several tools for training and using diffusion models.
%% Cell type:code id: tags:
``` python
fromdiffusersimportDDPMScheduler
# TO DO: Create the scheduler
noise_scheduler=
image,_=mnist_dataset[0]
# TO DO: Create a noise tensor sampled from a normal distribution with the same shape as the image
For the reverse diffusion process we will use a neural network. Given a noisy image and the corresponding timestep, the goal of the neural network is to predict the noise, which allows for the denoising.
For the model, we will have a similar architecture as we used for the cGAN generator, a 2D UNet with a few modifications. The main difference will be that we have to indicate to the model which timestep is currently being denoised. For that purpose, a timestep embedding is added, therefore the model has 2 inputs, the noisy image and the corresponding timestep.
In this exercise, we will use an UNet implementation from the diffusers library, which already has the timestep embedding included.
If the training takes too long, you can download an already trained model from [this link](https://partage.liris.cnrs.fr/index.php/s/AP2t6b3w8SM4Bp5) and use it for inference.
%% Cell type:code id: tags:
``` python
# TO DO: Add the path to the model checkpoint for loading the model
diffusion_backbone.load_state_dict(torch.load())
diffusion_backbone.eval()
```
%% Cell type:markdown id: tags:
### Time to generate some images.
During training, for each data sample, we take a random timestep and the corresponding noisy image to give it as input to our model. With sufficient training, the model should learn how to predict the noise in a noisy image for all possible timesteps.
During inference, to generate an image, we will start from pure noise and step by step predict the noise to go from one noisy image to the next, progressively denoising the image until we reach the timestep 0, in which we should have an image without any noise.
%% Cell type:code id: tags:
``` python
fromtqdmimporttqdm
# Start the image as random noise
image=torch.randn((10,1,64,64)).to(device)
# Create a list of images and labels for visualization
The diffusers library also provides *Pipeline* classes, which are wrappers around the model that abstracts the inference loop implemented above.
We can create a pipeline, giving it the trained model and the noise scheduler, and use it to generate images. In this case, we will only have access to the final image, generated on the last timestep, but not the intermediary images from the denoising process.
In this exercise we achieved decent results for very simple datasets. But we are quite far from those beautiful AI generated images we can find online. That is for 2 main reasons:
- Model size: due to the computation and time constrains, we can't really train very large models
- Dataset size: due to the same constrains, we can't use very complex and large datasets, which requires larger models and longer training times.
Fortunately, even though we can't train those large models with the available hardware and time, we can at least use them for inference !
The goal of this part is to learn how to retrieve and use a pre-trained diffusion models and also to get creative to come up with some nice prompts to generate outstanding images.
We are going to use Stable Diffusion 3.5, which is a state of the art open-source text conditioned model. It takes a prompt in natural language and use it to guide the diffusion process. This type of models are trained with image-text pairs, but can generalize beyond the pairs seen during training, being able to mix several different concepts into a single image.
In order to save memory, we will use quantization, which consists into converting the model weights types from float16 into float4. That simply means that each model weight will be stored using only 4 bits instead of 16 bits. That allows us to run the model in GPUs with less VRAM and have faster inference, with a small drop in the quality of the results.
%% Cell type:markdown id: tags:
For this part of the assignment restart the notebook kernel to be sure your GPU memory is empty. The memory usage can be verified using the command `nvidia-smi` in a terminal. If your GPU has 2GB of VRAM or less, the model will probably not fit into memory even with quantization. In that case, use Google Colab for this part or use the smaller model indicated below. If you are not happy with the results and have plenty of VRAM available, feel free to increase the quantization to 8 bits or even load the model without quantization.
# TO DO: test different prompts and visualize the generated images
# Once you are happy with the results, you can save 3 different images as png file with the correspondent prompts in a text file
prompt=
image=pipeline(
prompt=prompt,
num_inference_steps=40,
guidance_scale=4.5,
max_sequence_length=512
).images[0]
image.save("generated_image.png")
```
%% Cell type:markdown id: tags:
If even with the quantization you still run out of GPU memory and you can't use Google Colab, you can use the following code instead, which uses a much smaller model (the results won't be as near as impressive, but it should be able to run even on a CPU, if you have a little patience)
