10.3) (Exercise) Autoencoders, Generative Adversarial Networks, and Diffusion Models

import sys

# Is this notebook running on Colab or Kaggle?
IS_COLAB = "google.colab" in sys.modules
IS_KAGGLE = "kaggle_secrets" in sys.modules

import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from sklearn.manifold import TSNE

# make notebook reproducible
torch.manual_seed(42)

# make plot prettier
plt.rc('font', size=14)
plt.rc('axes', labelsize=14, titlesize=14)
plt.rc('legend', fontsize=14)
plt.rc('xtick', labelsize=10)
plt.rc('ytick', labelsize=10)

if not torch.cuda.is_available() and not torch.backends.mps.is_available():
    print("No GPU was detected. Training GANs and diffusion models can be very slow without a GPU.")
    if IS_COLAB:
        print("Go to Runtime > Change runtime and select a GPU hardware accelerator.")
    if IS_KAGGLE:
        print("Go to Settings > Accelerator and select GPU.")
    device = "cpu"
else:
    device = "cuda" if torch.cuda.is_available() else "mps"
    print(f"GPU runtime succesfully selected! We're ready to train our generative models.")

# for easy plotting later on
def plot_multiple_images(images, n_cols=None):
    n_cols = n_cols or len(images)
    n_rows = (len(images) - 1) // n_cols + 1
    if images.shape[-1] == 1:
        images = images.squeeze(axis=-1)
    plt.figure(figsize=(n_cols, n_rows))
    for index, image in enumerate(images):
        plt.subplot(n_rows, n_cols, index + 1)
        plt.imshow(image, cmap="binary")
        plt.axis("off")

def fit_reconstruction_model(model, loss_fn, optimizer, X_train, X_valid, epochs, batch_size=32):
    """Trains `model` to reconstruct its own input (X == y) -- the manual-training-loop
    equivalent of Keras's `model.fit(X_train, X_train, validation_data=(X_valid, X_valid))`."""
    train_tensor = torch.from_numpy(X_train).float().to(device)
    valid_tensor = torch.from_numpy(X_valid).float().to(device)
    dataset = torch.utils.data.TensorDataset(train_tensor)
    dataloader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=True)

    history = {"loss": [], "val_loss": []}
    for epoch in range(epochs):
        model.train()
        batch_losses = []
        for (X_batch,) in dataloader:
            optimizer.zero_grad()
            reconstructions = model(X_batch)
            loss = loss_fn(reconstructions, X_batch)
            loss.backward()
            optimizer.step()
            batch_losses.append(loss.item())

        model.eval()
        with torch.no_grad():
            val_loss = loss_fn(model(valid_tensor), valid_tensor).item()

        history["loss"].append(np.mean(batch_losses))
        history["val_loss"].append(val_loss)
        print(f"Epoch {epoch + 1}/{epochs} - loss: {history['loss'][-1]:.4f} - val_loss: {val_loss:.4f}")

    return history

# Load the CIFAR-10 dataset using torchvision.datasets.CIFAR10

# Normalize the pixel values to the range [0, 1]
# RGB values are between 0 and 256

# Split the training set into a training set and a validation set
# Get 5000 images as the validation set

# Load the CIFAR-10 dataset
train_set = torchvision.datasets.CIFAR10(root="_files/cifar10", train=True, download=True)
test_set = torchvision.datasets.CIFAR10(root="_files/cifar10", train=False, download=True)
X_train_full, y_train_full = train_set.data, np.array(train_set.targets)
X_test, y_test = test_set.data, np.array(test_set.targets)

# Normalize the pixel values to the range [0, 1]
# RGB values are between 0 and 256
X_train_full = X_train_full / __
X_test = X_test / __

# Split the training set into a training set and a validation set
# Get 5000 images as the validation set
X_train, X_valid = X_train_full[___], X_train_full[___]
y_train, y_valid = y_train_full[___], y_train_full[___]

### Get familiar with CIFAR-10 dataset
# Get the shape of the images

# Display a sample image

### Get familiar with CIFAR-10 dataset
# Get the shape of the images
print(___.shape, ___.shape)
# Display a sample image
____.____

# Define the stacked encoder architecture

# Define the stacked decoder architecture

# Combine encoder and decoder into the stacked autoencoder

# Compile the stacked autoencoder

# Train the stacked autoencoder
# We want the predictions to be as the input (X == y)

# Define the stacked encoder architecture
# Recommended 512, 256 units for the Dense layers and ReLU activation
stacked_encoder = nn.Sequential(
    nn.Flatten(),
    nn.LazyLinear(___), nn.___(),
    nn.LazyLinear(___), nn.___(),
)

# Define the stacked decoder architecture
# Recommended 512 and pixel count in one image as units for the Dense layers
stacked_decoder = nn.Sequential(
    nn.LazyLinear(___), nn.ReLU(),
    nn.LazyLinear(__ * __ * _),
    nn.Unflatten(1, (32, 32, 3)),
)

# Combine encoder and decoder into the stacked autoencoder
stacked_ae = nn.Sequential(___, ___)

# Compile the stacked autoencoder
# Recommended loss is MSE and recommended optimizer is NAdam
loss_fn = ___
optimizer = torch.optim.___(stacked_ae.parameters(), lr=___)

# Train the stacked autoencoder
# We want the predictions to be as the input (X == y)
# Recommended epochs is 10
history = fit_reconstruction_model(stacked_ae, loss_fn, optimizer, X_train, X_valid, epochs=__)

def plot_reconstructions(model, images=X_valid, n_images=5):
    model.eval()
    with torch.no_grad():
        inputs = torch.from_numpy(images[:n_images]).float().to(device)
        reconstructions = np.clip(model(inputs).cpu().numpy(), 0, 1)
    fig = plt.figure(figsize=(n_images * 1.5, 3))
    for image_index in range(n_images):
        plt.subplot(2, n_images, 1 + image_index)
        plt.imshow(images[image_index]) #, cmap="binary")
        plt.axis("off")
        plt.subplot(2, n_images, 1 + n_images + image_index)
        plt.imshow(reconstructions[image_index]) #, cmap="binary")
        plt.axis("off")

plot_reconstructions(stacked_ae)
plt.show()

# Predict the validation set

# Apply t-SNE for dimensionality reduction

# Transform the compressed validation set into 2D

# Plot the 2D data

# Predict the validation set
stacked_encoder.eval()
with torch.no_grad():
    X_valid_compressed = stacked_encoder(torch.from_numpy(___).float().to(device)).cpu().numpy()

# Apply t-SNE for dimensionality reduction
# You can initializes with PCA, and the learning_rate to auto
tsne = TSNE(init=___, learning_rate=___, random_state=42)

# Transform the compressed validation set into 2D
X_valid_2D = tsne.fit_transform(___)

# Plot the 2D data
plt.scatter(X_valid_2D[:, 0], X_valid_2D[:, 1], c=y_valid, s=10, cmap="tab10")
plt.show()

plt.figure(figsize=(10, 8))
cmap = plt.cm.tab10
Z = X_valid_2D
Z = (Z - Z.min()) / (Z.max() - Z.min())  # normalize to the 0-1 range
plt.scatter(Z[:, 0], Z[:, 1], c=y_valid, s=10, cmap=cmap)
image_positions = np.array([[1., 1.]])
for index, position in enumerate(Z):
    dist = ((position - image_positions) ** 2).sum(axis=1)
    if dist.min() > 0.02: # if far enough from other images
        image_positions = np.r_[image_positions, [position]]
        imagebox = mpl.offsetbox.AnnotationBbox(
            mpl.offsetbox.OffsetImage(X_valid[index], cmap="binary"),
            position, bboxprops={"edgecolor": cmap(y_valid[index]), "lw": 2})
        plt.gca().add_artist(imagebox)

plt.axis("off")
plt.show()

# PyTorch has no built-in equivalent to Keras's GaussianNoise layer -- like Dropout,
# it should only add noise during training, not at evaluation time
class GaussianNoise(nn.Module):
    def __init__(self, stddev):
        super().__init__()
        self.stddev = stddev

    def forward(self, x):
        if self.training:
            return x + torch.randn_like(x) * self.stddev
        return x

# Define the denoising encoder
denoising_encoder = nn.Sequential(
    # GaussianNoise adds noise for robustness (0.1) -- active only during training
    GaussianNoise(___),
    # Conv2d extracts 32 feature maps from the 3 RGB input channels, (3x3) kernel,
    # padding=1 to preserve size, ReLU as activation
    nn.Conv2d(___, ___, kernel_size=___, padding=___), nn.___(),
    nn.MaxPool2d(2),
    nn.Flatten(),
    # Dense layer with 512 units for feature processing, ReLU as activation
    nn.LazyLinear(___), nn.___(),
)

# Define the denoising decoder architecture
denoising_decoder = nn.Sequential(
    # Dense layer reshapes the compressed data back to match the decoder's input shape
    nn.LazyLinear(___ * ___ * ___), nn.___(),
    # Unflatten changes the 1D vector back into 32x16x16 feature maps (channels first)
    nn.Unflatten(1, (___, ___, ___)),
    # ConvTranspose2d performs upsampling (opposite of Conv2d) to restore the original image size.
    # Use 32 input channels, 3 output channels, 3 as kernel size, 2 as stride,
    # padding=1 and output_padding=1 (this pair doubles the spatial size, PyTorch's
    # equivalent of Keras's `padding="same"` for a stride-2 transposed convolution),
    # sigmoid as activation
    nn.ConvTranspose2d(___, ___, kernel_size=___, stride=___, padding=___, output_padding=___), nn.___(),
)

# Combine encoder and decoder into the denoising autoencoder
denoising_ae = nn.Sequential(___, ___)

# Compile the autoencoder
# Using binary crossentropy for the loss function and Nadam optimizer
loss_fn = ___
optimizer = torch.optim.___(denoising_ae.parameters())

# Train the autoencoder
# Input and target are the same (denoising task), with 10 epochs and validation data
history = fit_reconstruction_model(___, ___, ___, X_train, X_valid, epochs=___)

# Number of images to process (e.g. 5)
n_images = __

# Select a subset of test images
new_images = X_test[:___]

# Add noise to these images and scale it by various factors (e.g., 0.1)
new_images_noisy = new_images + np.random.randn(___, 32, 32, 3) * ___

# Predict denoised images using the autoencoder
denoising_ae.eval()
with torch.no_grad():
    new_images_denoised = denoising_ae(torch.from_numpy(new_images_noisy).float().to(device)).cpu().numpy()

# Plot the original, noisy and denoised images
plt.figure(figsize=(6, n_images * 2))
for index in range(n_images):
    plt.subplot(n_images, 3, index * 3 + 1)
    plt.imshow(new_images[index])
    plt.axis('off')
    if index == 0:
        plt.title("Original")
    plt.subplot(n_images, 3, index * 3 + 2)
    plt.imshow(new_images_noisy[index].clip(0., 1.))
    plt.axis('off')
    if index == 0:
        plt.title("Noisy")
    plt.subplot(n_images, 3, index * 3 + 3)
    plt.imshow(new_images_denoised[index])
    plt.axis('off')
    if index == 0:
        plt.title("Denoised")

plt.show()

# Define a custom PyTorch module for sampling from a normal distribution
class Sampling(nn.Module):
    def forward(self, mean, log_var):
        # Sample using the reparameterization trick
        return torch.randn_like(log_var) * torch.exp(log_var / 2) + mean

# Define the size of the latent space (e.g. 10)
codings_size = ___

class VariationalEncoder(nn.Module):
    def __init__(self, codings_size):
        super().__init__()
        self.flatten = nn.Flatten()
        # Pass the flattened 32x32x3 CIFAR image through Dense layers with ReLU activation
        self.dense1 = nn.LazyLinear(___)
        self.dense2 = nn.LazyLinear(___)
        # Compute mean and log variance for the latent space
        self.mean_layer = nn.LazyLinear(___)  # μ
        self.log_var_layer = nn.LazyLinear(___)  # γ
        self.sampling = Sampling()

    def forward(self, x):
        Z = self.flatten(x)
        Z = F.relu(self.dense1(Z))
        Z = F.relu(self.dense2(Z))
        codings_mean = self.mean_layer(Z)
        codings_log_var = self.log_var_layer(Z)
        # Sample from the latent space using the mean and log variance
        codings = self.sampling(___, ___)
        # Return mean, log variance and codings, like the reference's variational_encoder
        return codings_mean, codings_log_var, codings

variational_encoder = VariationalEncoder(codings_size)

class VariationalDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        # Recommended number of units are: 100, 150, and image size
        self.dense1 = nn.LazyLinear(___)
        self.dense2 = nn.LazyLinear(___)
        self.dense3 = nn.LazyLinear(___)

    def forward(self, codings):
        # Pass through Dense layers to reconstruct the original image
        x = F.relu(self.dense1(codings))
        x = F.relu(self.dense2(x))
        x = self.dense3(x)
        # Reshape as 32x32x3 CIFAR images
        return x.view(-1, ___, ___, ___)

variational_decoder = VariationalDecoder()

class VariationalAutoencoder(nn.Module):
    def __init__(self, encoder, decoder):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder

    def forward(self, x):
        # Encode inputs to get latent space codings
        codings_mean, codings_log_var, codings = self.encoder(___)
        # Decode codings to reconstruct the inputs
        reconstructions = self.decoder(___)
        return reconstructions, codings_mean, codings_log_var

variational_ae = VariationalAutoencoder(variational_encoder, variational_decoder)

def vae_latent_loss(codings_mean, codings_log_var):
    latent_loss = -0.5 * torch.sum(
        1 + codings_log_var - codings_log_var.exp() - codings_mean.pow(2),
        dim=-1)
    return latent_loss.mean() / 784.

# Compile the variational autoencoder
# Use Mean Squared Error for loss and Nadam optimizer for training
recon_loss_fn = ___
optimizer = torch.optim.___(variational_ae.parameters())

# Train the variational autoencoder
# Fit the model using training data with e.g. 25 epochs and e.g. 128 batch size
epochs = ___
batch_size = ___

train_tensor = torch.from_numpy(X_train).float().to(device)
valid_tensor = torch.from_numpy(X_valid).float().to(device)
dataloader = torch.utils.data.DataLoader(
    torch.utils.data.TensorDataset(train_tensor), batch_size=batch_size, shuffle=True)

for epoch in range(epochs):
    variational_ae.train()
    for (X_batch,) in dataloader:
        optimizer.zero_grad()
        reconstructions, codings_mean, codings_log_var = variational_ae(X_batch)
        loss = recon_loss_fn(reconstructions, X_batch) + vae_latent_loss(codings_mean, codings_log_var)
        loss.backward()
        optimizer.step()

    variational_ae.eval()
    with torch.no_grad():
        reconstructions, codings_mean, codings_log_var = variational_ae(valid_tensor)
        val_loss = recon_loss_fn(reconstructions, valid_tensor) + vae_latent_loss(codings_mean, codings_log_var)
    print(f"Epoch {epoch + 1}/{epochs} - val_loss: {val_loss.item():.4f}")

def plot_vae_reconstructions(model, images=X_valid, n_images=5):
    model.eval()
    with torch.no_grad():
        inputs = torch.from_numpy(images[:n_images]).float().to(device)
        reconstructions, _, _ = model(inputs)
        reconstructions = np.clip(reconstructions.cpu().numpy(), 0, 1)
    fig = plt.figure(figsize=(n_images * 1.5, 3))
    for image_index in range(n_images):
        plt.subplot(2, n_images, 1 + image_index)
        plt.imshow(images[image_index]) #, cmap="binary")
        plt.axis("off")
        plt.subplot(2, n_images, 1 + n_images + image_index)
        plt.imshow(reconstructions[image_index]) #, cmap="binary")
        plt.axis("off")

plot_vae_reconstructions(variational_ae)
plt.show()

# Define the size of the latent space

# Build the generator model

# Build the discriminator model

# PyTorch doesn't need a combined "GAN" model or a frozen-discriminator flag like Keras:
# calling discriminator(generator(noise)) directly and stepping only the generator's own
# optimizer already leaves the discriminator's weights untouched during that phase.

# Define the size of the latent space, e.g. 30
codings_size = ___

# Build the generator model
generator = nn.Sequential(
    nn.LazyLinear(___), nn.___(),  # Expand to 300 units
    nn.LazyLinear(___), nn.___(),  # Expand to 450 units
    nn.LazyLinear(___ * ___ * ___), nn.Sigmoid(),  # Output layer to match 32x32x3 image
    nn.Unflatten(1, (___, ___, ___)),  # Reshape to 32x32x3 CIFAR image
)

# Build the discriminator model
discriminator = nn.Sequential(
    nn.Flatten(),  # Flatten the input image
    nn.LazyLinear(___), nn.___(),  # Hidden layer with 450 units
    nn.LazyLinear(___), nn.___(),  # Hidden layer with 300 units
    nn.LazyLinear(1), nn.Sigmoid(),  # Output layer for binary classification
)

# Compile the discriminator model

# Set discriminator to non-trainable when training the GAN to freeze its weights

# Compile the GAN model

# Set up the discriminator's loss and optimizer
# Uses binary cross-entropy loss for binary classification and RMSprop optimizer
loss_fn = ___
disc_optimizer = torch.optim.___(discriminator.parameters(), lr=___)

# A separate optimizer holding only the generator's parameters -- this is what keeps the
# discriminator "frozen" while training the generator, PyTorch's equivalent of Keras's
# `discriminator.trainable = False`
gan_optimizer = torch.optim.___(generator.parameters(), lr=___)

# Define batch size for training

# Create a PyTorch dataset from the training data

# Batch the data into chunks of size batch_size

# Define batch size for training
batch_size = ___

# Create a torch Dataset from the training data
dataset = torch.utils.data.TensorDataset(torch.from_numpy(X_train).float())

# Shuffle and batch the data -- DataLoader's num_workers can prefetch batches in the
# background, PyTorch's equivalent of tf.data's .prefetch()
dataloader = torch.utils.data.DataLoader(dataset, batch_size=___, shuffle=True, drop_last=True)

# Helper function to train the GAN
def train_gan(generator, discriminator, disc_optimizer, gan_optimizer, loss_fn,
              dataloader, codings_size, n_epochs):
    for epoch in range(n_epochs):
        print(f"Epoch {epoch + 1}/{n_epochs}")

        for (X_batch,) in dataloader:
            X_batch = X_batch.to(device)
            batch_size = X_batch.shape[0]

            # phase 1 - training the discriminator
            noise = torch.randn(batch_size, codings_size, device=device)
            generated_images = generator(noise)
            X_fake_and_real = torch.cat([generated_images.detach(), X_batch], dim=0)
            y1 = torch.cat([torch.zeros(batch_size, 1), torch.ones(batch_size, 1)]).to(device)

            discriminator.train()
            disc_optimizer.zero_grad()
            loss = loss_fn(discriminator(X_fake_and_real), y1)
            loss.backward()
            disc_optimizer.step()

            # phase 2 - training the generator
            noise = torch.randn(batch_size, codings_size, device=device)
            y2 = torch.ones(batch_size, 1, device=device)

            gan_optimizer.zero_grad()
            loss = loss_fn(discriminator(generator(noise)), y2)
            loss.backward()
            gan_optimizer.step()

        plot_multiple_images(generated_images.detach().cpu().numpy(), 8)
        plt.show()

# Train the GAN model, with e.g. 50 epochs

# Train the GAN model, with e.g. 50 epochs
train_gan(generator, discriminator, disc_optimizer, gan_optimizer, loss_fn,
          dataloader, codings_size, n_epochs=___)

# Generate a batch of latent vectors

# Generate images using the trained generator (using the `codings`)

# Plot the generated images, e.g. 5

# Generate a batch of latent vectors
codings = torch.randn(batch_size, codings_size, device=device)

# Generate images using the trained generator (using the `codings`)
generator.eval()
with torch.no_grad():
    generated_images = generator(codings).cpu().numpy()

# Plot the generated images, e.g. 5
plot_multiple_images(generated_images, ___)
plt.show()

# Define the size of the latent space, e.g. 100
codings_size = ___

# Build the generator model
# Generates images from the latent space vector
# First Dense layer expands to 8x8x128, then to be reshaped (channels first: 128x8x8)
generator = nn.Sequential(
    nn.LazyLinear(___ * ___ * ___),  # Expand to 8x8x128
    nn.Unflatten(1, (___, ___, ___)),  # Reshape to 128x8x8
    nn.BatchNorm2d(___),
    # Upsample to 64 channels, 5 as kernel size, 2 strides, padding=2, output_padding=1
    # (PyTorch's equivalent of Keras's `padding="same"` for a stride-2 transposed convolution)
    nn.ConvTranspose2d(___, ___, kernel_size=___, stride=___,
                       padding=___, output_padding=___), nn.ReLU(),
    nn.BatchNorm2d(___),
    # Output layer with 3 channels, 5 as kernel size, 2 strides, padding=2, output_padding=1
    nn.ConvTranspose2d(___, ___, kernel_size=___, stride=___,
                       padding=___, output_padding=___), nn.Tanh(),
)

# Build the discriminator model that classifies images as real or fake
discriminator = nn.Sequential(
    # Downsample to 64; 5 as kernel size, 2 strides, padding=2
    nn.Conv2d(___, ___, kernel_size=___, stride=___, padding=___), nn.LeakyReLU(0.2),
    nn.Dropout(___),  # e.g. 0.4
    # Downsample to 128; 5 as kernel size, 2 strides, padding=2
    nn.Conv2d(___, ___, kernel_size=___, stride=___, padding=___), nn.LeakyReLU(0.2),
    nn.Dropout(___),  # e.g. 0.4
    nn.Flatten(),
    nn.LazyLinear(1), nn.Sigmoid(),
)

# Set up the discriminator's loss and optimizer, binary cross-entropy and RMSprop
loss_fn = ___
disc_optimizer = torch.optim.___(discriminator.parameters())

# A separate optimizer for the generator only -- see the note above cell 71
gan_optimizer = torch.optim.___(generator.parameters())

# Reshape to channels-first (N, 3, 32, 32) and scale to match the generator's [-1, 1]
# output range (nn.Tanh)
X_train_dcgan = X_train.transpose(0, 3, 1, 2) * 2. - 1.

# Set the batch size for training
batch_size = ___

# Create a dataset from reshaped and rescaled training data
dataset = torch.utils.data.TensorDataset(torch.from_numpy(X_train_dcgan).float())

# Shuffle and batch the dataset
dataloader = torch.utils.data.DataLoader(dataset, batch_size=___, shuffle=True, drop_last=True)

# Train the GAN model with e.g. 50 epochs
train_gan(generator, discriminator, disc_optimizer, gan_optimizer, loss_fn,
          dataloader, codings_size, n_epochs=___)

# Generate random noise for input to the generator
# `batch_size` is the number of samples, `codings_size` is the latent space size
noise = torch.randn(___, ___, device=device)

# Generate images using the generator model
generator.eval()
with torch.no_grad():
    generated_images = generator(noise).cpu()

# Plot e.g. 5 generated images -- move channels back to last, rescale from [-1,1] to [0,1]
generated_images = (generated_images.permute(0, 2, 3, 1).numpy() + 1) / 2
plot_multiple_images(generated_images, ___)

def variance_schedule(T, s=0.008, max_beta=0.999):
    t = np.arange(T + 1)
    f = np.cos((t / T + s) / (1 + s) * np.pi / 2) ** 2
    alpha = np.clip(f[1:] / f[:-1], 1 - max_beta, 1)
    alpha = np.append(1, alpha).astype(np.float32)  # add α₀ = 1
    beta = 1 - alpha
    alpha_cumprod = np.cumprod(alpha)
    return alpha, alpha_cumprod, beta  # αₜ , α̅ₜ , βₜ for t = 0 to T

np.random.seed(42)  # extra code – for reproducibility
T = 4000
alpha, alpha_cumprod, beta = variance_schedule(T)

plt.figure(figsize=(6, 3))
plt.plot(beta, "r--", label=r"$\beta_t$")
plt.plot(alpha_cumprod, "b", label=r"$\bar{\alpha}_t$")
plt.axis([0, T, 0, 1])
plt.grid(True)
plt.xlabel(r"t")
plt.legend()
plt.show()

# Move the variance schedule to torch tensors, so it can be indexed with a batch of
# timesteps directly (prepare_batch and generate, below, both need this)
alpha = torch.from_numpy(alpha).to(device)
alpha_cumprod = torch.from_numpy(alpha_cumprod.astype(np.float32)).to(device)
beta = torch.from_numpy(beta).to(device)

def prepare_batch(X):
    X = X.permute(0, 3, 1, 2).float() * 2 - 1  # channels-last -> channels-first, scale to [-1, +1]
    batch_size = X.shape[0]
    t = torch.randint(1, T + 1, (batch_size,), device=X.device)
    alpha_cm = alpha_cumprod[t].view(batch_size, 1, 1, 1)
    noise = torch.randn_like(X)
    return {
        "X_noisy": alpha_cm ** 0.5 * X + (1 - alpha_cm) ** 0.5 * noise,
        "time": t,
    }, noise

def prepare_dataset(X, batch_size=32, shuffle=False):
    # Create a dataset from input data
    dataset = torch.utils.data.TensorDataset(torch.from_numpy(X))

    # Return a DataLoader; shuffling and batching happen here, prepare_batch() itself
    # is called explicitly at each training step (see the training loop below)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

# Prepare the training and validation datasets

def prepare_dataset(X, batch_size=32, shuffle=False):
    # Create a dataset from input data
    dataset = torch.utils.data.TensorDataset(torch.from_numpy(X))

    # Return a DataLoader; shuffling and batching happen here, prepare_batch() itself
    # is called explicitly at each training step (see the training loop below)
    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle)

# Prepare the training and validation datasets
train_loader = prepare_dataset(X_train, batch_size=___, shuffle=___)
valid_loader = prepare_dataset(X_valid, batch_size=___)

def subtract_noise(X_noisy, time, noise):
    alpha_cm = alpha_cumprod[time].view(-1, 1, 1, 1)
    return (X_noisy - (1 - alpha_cm) ** 0.5 * noise) / alpha_cm ** 0.5

(X_batch,) = next(iter(train_loader))  # get the first batch
X_dict, Y_noise = prepare_batch(X_batch)
X_original = subtract_noise(X_dict["X_noisy"], X_dict["time"], Y_noise)

# Plot original images, noisy images and the noise to predict
def to_image_grid(X):
    return ((X.permute(0, 2, 3, 1).detach().cpu().numpy() + 1) * 128).clip(0, 255).astype(np.uint8)

print("Original images")
plot_multiple_images(to_image_grid(X_original[:8]))
plt.show()
print("Time steps:", X_dict["time"].cpu().numpy()[:8])
print("Noisy images")
plot_multiple_images(to_image_grid(X_dict["X_noisy"][:8]))
plt.show()
print("Noise to predict")
plot_multiple_images(to_image_grid(Y_noise[:8]))
plt.show()

embed_size = 64

class TimeEncoding(nn.Module):
    def __init__(self, T, embed_size):
        super().__init__()
        assert embed_size % 2 == 0, "embed_size must be even"
        p, i = np.meshgrid(np.arange(T + 1), 2 * np.arange(embed_size // 2))
        t_emb = np.empty((T + 1, embed_size))
        t_emb[:, ::2] = np.sin(p / 10_000 ** (i / embed_size)).T
        t_emb[:, 1::2] = np.cos(p / 10_000 ** (i / embed_size)).T
        self.register_buffer("time_encodings", torch.tensor(t_emb, dtype=torch.float32))

    def forward(self, t):
        return self.time_encodings[t]

# the size of the embedding e.g. 64
embed_size = ___

# Custom module to encode time steps with sinusoidal embeddings
class TimeEncoding(nn.Module):
    def __init__(self, T, embed_size):
        # Initialize the module and ensure embed_size is even
        super().__init__()
        assert embed_size % 2 == 0, "embed_size must be even"

        # Create a meshgrid for time steps and embedding indices
        p, i = np.meshgrid(np.arange(T + 1), 2 * np.arange(embed_size // 2))

        # Initialize the time embeddings matrix
        t_emb = np.empty((T + 1, embed_size))

        # Fill even indices with sine values and odd indices with cosine values
        t_emb[:, ::2] = np.sin(p / 10_000 ** (i / embed_size)).T
        t_emb[:, 1::2] = np.cos(p / 10_000 ** (i / embed_size)).T

        # Register the embeddings as a (non-trainable) buffer
        self.register_buffer("time_encodings", torch.tensor(t_emb, dtype=torch.float32))

    # Method to fetch time encodings for `t` time steps
    def forward(self, t):
        return self.time_encodings[___]

# the size of the embedding e.g. 64
embed_size = ___

# Custom module to encode time steps with sinusoidal embeddings
class TimeEncoding(nn.Module):
    def __init__(self, T, embed_size):
        # Initialize the module and ensure embed_size is even
        super().__init__()
        assert embed_size % 2 == 0, "embed_size must be even"

        # Create a meshgrid for time steps and embedding indices
        p, i = np.meshgrid(np.arange(T + 1), 2 * np.arange(embed_size // 2))

        # Initialize the time embeddings matrix
        t_emb = np.empty((T + 1, embed_size))

        # Fill even indices with sine values and odd indices with cosine values
        t_emb[:, ::2] = np.sin(p / 10_000 ** (i / embed_size)).T
        t_emb[:, 1::2] = np.cos(p / 10_000 ** (i / embed_size)).T

        # Register the embeddings as a (non-trainable) buffer
        self.register_buffer("time_encodings", torch.tensor(t_emb, dtype=torch.float32))

    # Method to fetch time encodings for `t` time steps
    def forward(self, t):
        return self.time_encodings[t]

class DiffusionModel(nn.Module):
    """A UNet-like model: Conv2d + BatchNorm2d blocks and skip connections, gradually
    downsampling the image (stride-2 pooling), then growing it back again (nearest-neighbor
    upsampling). Skip connections are added across the downsampling and upsampling parts.
    Time encodings are added to the output of each block, after resizing them with a Linear
    layer.

    Note: an image's time encoding is added to every pixel in the image, along the channel
    axis -- so the number of channels going into each block must match the embedding
    projection's output size.

    Unlike the reference this is adapted from (which was built for 28x28x1 Fashion MNIST),
    this version works directly on CIFAR-10's native 32x32x3 shape -- no zero-padding /
    cropping round-trip is needed since 32 is already evenly divisible by 2 three times over.
    """
    def __init__(self, T, embed_size=64, dim=16):
        super().__init__()
        self.time_encoding = TimeEncoding(T, embed_size)

        self.init_conv = nn.Conv2d(3, dim, kernel_size=3, padding=1)
        self.init_bn = nn.BatchNorm2d(dim)
        self.init_time = nn.Linear(embed_size, dim)

        down_dims = (32, 64, 128)
        in_dim = dim
        self.down_convs = nn.ModuleList()
        self.down_skips = nn.ModuleList()
        self.down_times = nn.ModuleList()
        for d in down_dims:
            self.down_convs.append(nn.ModuleList([
                nn.Conv2d(in_dim, d, kernel_size=3, padding=1),
                nn.BatchNorm2d(d),
                nn.Conv2d(d, d, kernel_size=3, padding=1),
                nn.BatchNorm2d(d),
            ]))
            self.down_skips.append(nn.Conv2d(in_dim, d, kernel_size=1, stride=2))
            self.down_times.append(nn.Linear(embed_size, d))
            in_dim = d

        up_dims = (64, 32, 16)
        cross_dims = down_dims[::-1]
        self.up_convs = nn.ModuleList()
        self.up_skiplinks = nn.ModuleList()
        self.up_times = nn.ModuleList()
        for d, cross_d in zip(up_dims, cross_dims):
            self.up_convs.append(nn.ModuleList([
                nn.ConvTranspose2d(in_dim, d, kernel_size=3, padding=1),
                nn.BatchNorm2d(d),
                nn.ConvTranspose2d(d, d, kernel_size=3, padding=1),
                nn.BatchNorm2d(d),
            ]))
            self.up_skiplinks.append(nn.Conv2d(in_dim, d, kernel_size=1))
            self.up_times.append(nn.Linear(embed_size, d))
            in_dim = d + cross_d  # after concatenating with the matching downsampling skip

        self.final_conv = nn.Conv2d(in_dim, 3, kernel_size=3, padding=1)

    def forward(self, X_noisy, time):
        # Encode the time step using the TimeEncoding module
        time_enc = self.time_encoding(time)

        # Initial convolution
        Z = F.relu(self.init_bn(self.init_conv(X_noisy)))
        # Adapt the time encoding and add it to the image feature map
        t = self.init_time(time_enc)[:, :, None, None]
        Z = Z + t

        # Keep track of skip connections and initiate a residual connection
        skip = Z
        cross_skips = []  # for skip connections in the UNet structure

        # Downsampling blocks
        for (conv1, bn1, conv2, bn2), skip_conv, time_lin in zip(
                self.down_convs, self.down_skips, self.down_times):
            Z = F.relu(bn1(conv1(Z)))
            Z = F.relu(bn2(conv2(Z)))

            # Store intermediate output for skip connection
            cross_skips.append(Z)

            # Downsample and add residual connection
            Z = F.max_pool2d(Z, kernel_size=3, stride=2, padding=1)
            skip_link = skip_conv(skip)
            Z = Z + skip_link

            # Add time information to downsampled feature maps
            t = time_lin(time_enc)[:, :, None, None]
            Z = Z + t
            skip = Z

        # Upsampling blocks
        for (convT1, bn1, convT2, bn2), skip_conv, time_lin in zip(
                self.up_convs, self.up_skiplinks, self.up_times):
            Z = F.relu(bn1(convT1(Z)))
            Z = F.relu(bn2(convT2(Z)))

            # Upsample and add residual connection
            Z = F.interpolate(Z, scale_factor=2, mode="nearest")
            skip_link = F.interpolate(skip, scale_factor=2, mode="nearest")
            skip_link = skip_conv(skip_link)
            Z = Z + skip_link

            # Add time encoding and merge with the corresponding downsampling skip connection
            t = time_lin(time_enc)[:, :, None, None]
            Z = Z + t
            Z = torch.cat([Z, cross_skips.pop()], dim=1)
            skip = Z

        # Final convolution layer
        return self.final_conv(Z)

# Build and compile the diffusion model
model = DiffusionModel(T).to(device)
loss_fn = nn.HuberLoss()
optimizer = torch.optim.NAdam(model.parameters())

# Train the model with the training and validation datasets with e.g. 100 epochs,
# saving the best model as we go (the manual-training-loop equivalent of Keras's
# `ModelCheckpoint(save_best_only=True)`)
epochs = ___
best_val_loss = float("inf")
for epoch in range(epochs):
    model.train()
    train_losses = []
    for (X_batch,) in train_loader:
        X_dict, noise = prepare_batch(X_batch.to(device))
        optimizer.zero_grad()
        pred_noise = model(X_dict["X_noisy"], X_dict["time"])
        loss = loss_fn(pred_noise, noise)
        loss.backward()
        optimizer.step()
        train_losses.append(loss.item())

    model.eval()
    val_losses = []
    with torch.no_grad():
        for (X_batch,) in valid_loader:
            X_dict, noise = prepare_batch(X_batch.to(device))
            pred_noise = model(X_dict["X_noisy"], X_dict["time"])
            val_losses.append(loss_fn(pred_noise, noise).item())

    val_loss = np.mean(val_losses)
    print(f"Epoch {epoch + 1}/{epochs} - loss: {np.mean(train_losses):.4f} - val_loss: {val_loss:.4f}")

    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), "my_diffusion_model.pth")

model.load_state_dict(torch.load("my_diffusion_model.pth"))

def generate(model, batch_size=32):
    model.eval()
    X = torch.randn(batch_size, 3, 32, 32, device=device)
    with torch.no_grad():
        for t in range(T - 1, 0, -1):
            print(f"\rt = {t}", end=" ")  # show progress
            noise = torch.randn_like(X) if t > 1 else torch.zeros_like(X)
            time_batch = torch.full((batch_size,), t, dtype=torch.long, device=device)
            X_noise = model(X, time_batch)
            X = (
                1 / alpha[t] ** 0.5
                * (X - beta[t] / (1 - alpha_cumprod[t]) ** 0.5 * X_noise)
                + (1 - alpha[t]) ** 0.5 * noise
            )
    return X

# Generate images
X_gen = generate(model)

# Plot the generated images -- move channels back to last and rescale to [0, 1] for imshow
X_gen_grid = ((X_gen.permute(0, 2, 3, 1).cpu().numpy() + 1) / 2).clip(0, 1)
plot_multiple_images(X_gen_grid, 5)
plt.show()

def generate_ddim(model, batch_size=32, n_steps=50):
    """Sample with DDIM: the trained model, evaluated on a subsequence of timesteps."""
    model.eval()
    X = torch.randn(batch_size, 3, 32, 32, device=device)

    # a decreasing subsequence of the T training timesteps, e.g. 4000, 3918, ..., 1
    times = torch.linspace(T, 1, n_steps).long().tolist()

    with torch.no_grad():
        # pair each timestep with the one it jumps to; the last jump lands on 0, and
        # alpha_cumprod[0] is 1, which is the clean image
        for t, t_prev in zip(times, times[1:] + [0]):
            time_batch = torch.full((batch_size,), t, dtype=torch.long, device=device)
            X_noise = model(X, time_batch)

            # invert the forward process to get the clean image the model implies,
            # then re-noise it to t_prev rather than stepping through every timestep
            X_start = (X - (1 - alpha_cumprod[t]) ** 0.5 * X_noise) / alpha_cumprod[t] ** 0.5
            X_start = X_start.clamp(-1, 1)          # the estimate can drift outside the data range
            X = alpha_cumprod[t_prev] ** 0.5 * X_start + (1 - alpha_cumprod[t_prev]) ** 0.5 * X_noise

    return X

# Generate with the same trained model, timing both samplers to see the difference
import time

start = time.time()
X_gen_ddim = generate_ddim(model, n_steps=___)
print(f"DDIM: {time.time() - start:.1f} s")

X_gen_ddim_grid = ((X_gen_ddim.permute(0, 2, 3, 1).cpu().numpy() + 1) / 2).clip(0, 1)
plot_multiple_images(X_gen_ddim_grid, 5)
plt.show()