---
title: "Generative Deep Learning: Generating images with variational autoencoders"
output: 
  html_notebook: 
    theme: cerulean
    highlight: textmate
editor_options: 
  chunk_output_type: console
---

## 12.4 Generating images with variational autoencoders

```{r}
library(keras)
library(tensorflow)
library(tfdatasets)
library(tfautograph)
```

### 12.4.1 Sampling from latent spaces of images

The core idea of image generation is to develop a **low-dimensional latent space** (a vector space) where each point can be mapped to a **"valid" image**. A module called a **generator (for GANs) or decoder (for VAEs)** performs this mapping. Once this latent space is learned, new images can be generated by **sampling points from it** and mapping them back to the image space. These new images are essentially interpolations of the training images.

### 12.4.2 Concept vectors for image editing

In a learned **latent space** of image representations, certain **directions** can correspond to meaningful **axes of variation** in the original data. These directions are called **concept vectors**. By identifying such vectors (e.g., a "smile vector" in a face latent space), it becomes possible to **edit images** by projecting them into the latent space, moving their representation along the concept vector, and then decoding the modified representation back into an image. This allows for intuitive image manipulation like adding features (e.g., sunglasses) or changing attributes (e.g., from frowning to smiling).

### 12.4.3 Variational autoencoders

Variational autoencoders (VAEs) are generative models particularly useful for **image editing via concept vectors**. They represent a modern evolution of autoencoders, incorporating ideas from deep learning and Bayesian inference. Instead of a classical autoencoder compressing an input to a fixed latent code, a VAE encodes an input image into the **parameters of a statistical distribution**: a mean and a variance. This implies that the input image is assumed to have been generated by a statistical process. The VAE then uses these parameters to **randomly sample a point from this latent distribution** and decodes that point back into an image. This **stochasticity** forces the latent space to be **continuous and highly structured**, where every point in the latent space decodes to a valid output. The training of a VAE involves **two loss functions**: a **reconstruction loss** to ensure the decoded images match the inputs, and a **regularization loss** (typically the Kullback-Leibler divergence) to encourage the latent distribution to be a well-behaved normal distribution around zero.

The core steps in a **Variational Autoencoder (VAE)**: First, an **encoder** network maps an `input_img` to the **parameters of a latent probability distribution**: `z_mean` and `z_log_variance`. Then, a latent point `z` is **randomly sampled** from this distribution using the formula `z = z_mean + exp(z_log_variance) * epsilon`, where `epsilon` is a random tensor. Finally, a **decoder** network maps this sampled latent point `z` back to a `reconstructed_img`. The `keras_model` then defines the VAE model that maps the `input_img` to its `reconstructed_img`. This process allows the VAE to learn a **continuous and structured latent space** of the input data.

```{r eval=FALSE}
# Encode the input into mean and variance parameters. 
c(z_mean, z_log_variance) %<-% encoder(input_img)
# Draw a latent point using a small random epsilon.
z <- z_mean + exp(z_log_variance) * epsilon
# Decode z back to an image.
reconstructed_img <- decoder(z)
# Instantiate the autoencoder model, which maps an input image to its reconstruction.
model <- keras_model(input_img, reconstructed_img)
```

### 12.4.4 Implementing a VAE with Keras

We define the **encoder network** for a Variational Autoencoder (VAE) designed for MNIST digits. Its main purpose is to take an input **28x28 grayscale image** and map it to the **parameters of a 2-dimensional latent probability distribution**: the **mean (`z_mean`)** and the **log-variance (`z_log_var`)**.

The encoder consists of the following key layers:
*   An **input layer** expecting the image.
*   **Two 2D convolutional layers** with ReLU activation and a **stride of 2** for downsampling, increasing the number of filters from 32 to 64. The use of strides is preferred over max pooling for preserving location information.
*   A **flatten layer** to convert the convolutional output into a 1D tensor.
*   A **dense layer with 16 units** and ReLU activation.
*   **Two final dense layers**, each with 2 units (matching the `latent_dim`): one to output the **`z_mean` vector**, and the other to output the **`z_log_var` vector**.

The encoder model takes the input image and outputs a **list containing both the `z_mean` and `z_log_var`**, which will then be used by the sampling layer to generate a point in the latent space.

```{r}
# Dimensionality of the latent space: a 2D plane
latent_dim <- 2

encoder_inputs <-  layer_input(shape=c(28, 28, 1))
x <- encoder_inputs %>%
  layer_conv_2d(32, 3, activation = "relu", strides = 2, padding = "same") %>%
  layer_conv_2d(64, 3, activation = "relu", strides = 2, padding = "same") %>%
  layer_flatten() %>%
  layer_dense(16, activation = "relu")

# The input image ends up being encoded into these two parameters.
z_mean    <- x %>% layer_dense(latent_dim, name="z_mean")
z_log_var <- x %>% layer_dense(latent_dim, name="z_log_var")
encoder <- keras_model(encoder_inputs, list(z_mean, z_log_var),
                       name="encoder")

encoder
```

Latent-space-sampling layer

```{r}
layer_sampler <- new_layer_class(
  classname = "Sampler",
  # z_mean and z_log_var here both will have shape (batch_size, latent_dim).
  call = function(self, z_mean, z_log_var) {
    # Draw a batch of random normal vectors.
    epsilon <- tf$random$normal(shape = tf$shape(z_mean))
    # Apply the VAE sampling formula.
    z_mean + exp(0.5 * z_log_var) * epsilon
  }
)
```

VAE decoder network, mapping latent space points to images

```{r}
# Input where we’ll feed z 
latent_inputs <- layer_input(shape = c(latent_dim))
decoder_outputs <- latent_inputs %>%
  # Produce the same number of coefficients that we had at the
  # level of the Flatten layer in the encoder.
  layer_dense(7 * 7 * 64, activation = "relu") %>%
  # Revert the layer_flatten() of the encoder.
  layer_reshape(c(7, 7, 64)) %>%
  # Revert the layer_conv_2d() of the encoder.
  layer_conv_2d_transpose(64, 3, activation = "relu",
                          strides = 2, padding = "same") %>%
  layer_conv_2d_transpose(32, 3, activation = "relu",
                          strides = 2, padding = "same") %>%
  # The output ends up with shape (28, 28, 1).
  layer_conv_2d(1, 3, activation = "sigmoid", padding = "same")
decoder <- keras_model(latent_inputs, decoder_outputs,
                       name = "decoder")

decoder
```

VAE model with custom train_step()

```{r}
model_vae <- new_model_class(
  classname = "VAE",

  initialize = function(encoder, decoder, ...) {
    super$initialize(...)
    # We assign to self instead of private because self$encoder <- encoder
    # we want the layer weights automatically tracked by the Keras Model base class.
    self$encoder <- encoder
    self$decoder <- decoder
    self$sampler <- layer_sampler()
    # We use these metrics to keep track of the loss averages over each epoch.
    self$total_loss_tracker <-
      metric_mean(name = "total_loss")
    self$reconstruction_loss_tracker <-
      metric_mean(name = "reconstruction_loss")
    self$kl_loss_tracker <-
      metric_mean(name = "kl_loss")
  },

  # We list the metrics in an active property to enable the framework to reset them after
  # each epoch (or between multiple calls to fit()/evaluate()).
  metrics = mark_active(function() {
    list(
      self$total_loss_tracker,
      self$reconstruction_loss_tracker,
      self$kl_loss_tracker
    )
  }),

  train_step = function(data) {
    with(tf$GradientTape() %as% tape, {

      c(z_mean, z_log_var) %<-% self$encoder(data)
      z <- self$sampler(z_mean, z_log_var)

      reconstruction <- decoder(z)
      # We sum the reconstruction loss over the spatial dimensions (second and
      # reconstruction <- decoder(z) third axes) and take its mean over the batch dimension.
      reconstruction_loss <-
        loss_binary_crossentropy(data, reconstruction) %>%
        # Total loss for each case in the batch; preserve batch axis.
        sum(axis = c(2, 3)) %>%
        # Take the mean of loss totals in the batch.
        mean()

      kl_loss <- -0.5 * (1 + z_log_var - z_mean^2 - exp(z_log_var))
      # Add the regularization term (Kullback–Leibler divergence).
      total_loss <- reconstruction_loss + mean(kl_loss)
    })

    grads <- tape$gradient(total_loss, self$trainable_weights)
    self$optimizer$apply_gradients(zip_lists(grads, self$trainable_weights))

    self$total_loss_tracker$update_state(total_loss)
    self$reconstruction_loss_tracker$update_state(reconstruction_loss)
    self$kl_loss_tracker$update_state(kl_loss)

    list(total_loss = self$total_loss_tracker$result(),
         reconstruction_loss = self$reconstruction_loss_tracker$result(),
         kl_loss = self$kl_loss_tracker$result())
  }
)
```

Finally, we’re ready to instantiate and train the model on MNIST digits. Because the
loss is taken care of in the custom layer, we don’t specify an external loss at compile
time (loss = NULL), which in turn means we won’t pass target data during training (as
you can see, we pass only x_train to the model in fit()).

```{r}
# listarrays: A Toolbox for Working with R Arrays in a Functional Programming Style
library(listarrays) # https://CRAN.R-project.org/package=listarrays
c(c(x_train, .), c(x_test, .)) %<-% dataset_mnist()

mnist_digits <-
  # We train on all MNIST digits, so we combine the training and test samples along the batch dim.
  bind_on_rows(x_train, x_test) %>%
  expand_dims(-1) %>%
  { . / 255 }

str(mnist_digits)

vae <- model_vae(encoder, decoder)
# Note that we don’t pass a loss argument in compile(), because
# the loss is already part of the train_step().
vae %>% compile(optimizer = optimizer_adam())

# Note that we don’t pass targets in fit(), because train_step() doesn’t expect any.
# vae %>% fit(mnist_digits, epochs = 30, batch_size = 128)

```

About 30 min on CPU

```{r}
# Measure training time
start_time <- Sys.time()

# Fit the model
vae %>% fit(mnist_digits, epochs = 30, batch_size = 128)

# Calculate total training time
end_time <- Sys.time()
total_time <- end_time - start_time
cat(sprintf("✅ Total training time: %.2f seconds\n", as.numeric(total_time, units = "secs")))

vae %>% save_model_weights_tf("vae_weights")
cat("✅ Model weights saved successfully!\n")

# Recreate the original VAE model
vae_new <- model_vae(encoder, decoder)  # Rebuild model
vae_new %>% compile(optimizer = optimizer_adam())  # Recompile

# Load saved weights
vae_new %>% load_model_weights_tf("vae_weights")
cat("✅ Model weights loaded successfully!\n")
```

Sampling a grid of images from the 2D latent space

```{r}
n <- 30
digit_size <- 28

# Create a 2D grid of linearly spaced samples.
z_grid <-
  seq(-1, 1, length.out = n) %>%
  expand.grid(., .) %>%
  as.matrix()

# Get the decoded digits. 
decoded <- predict(vae$decoder, z_grid)

# Transform the decoded digits with shape (900, 28, 28, 1) to an R array with shape (28*30, 28*30) for plotting.
z_grid_i <- seq(n) %>% expand.grid(x = ., y = .)
# We’ll display a grid of 30 × 30 digits (900 digits total).
figure <- array(0, c(digit_size * n, digit_size * n))
for (i in 1:nrow(z_grid_i)) {
  c(xi, yi) %<-% z_grid_i[i, ]
  digit <- decoded[i, , , ]
  figure[seq(to = (n + 1 - xi) * digit_size, length.out = digit_size),
         seq(to = yi * digit_size, length.out = digit_size)] <-
    digit
}

# Square plot type
par(pty = "s")
# Expand lim so (–1, 1) are at the center of a digit.
lim <- extendrange(r = c(-1, 1),
                   f = 1 - (n / (n+.5)))
plot(NULL, frame.plot = FALSE,
     ylim = lim, xlim = lim,
     # Pass a formula object to xlab for a proper subscript.
     xlab = ~z[1], ylab = ~z[2])
# Subtract from 1 to invert the colors.
rasterImage(as.raster(1 - figure, max = 1),
            lim[1], lim[1], lim[2], lim[2],
            interpolate = FALSE)
```

