Now, we can extend our VAE architecture by adding another input to the decoder and encoder.
We will pass a class label in one-hot encoded format to let the model know which particular digit we are generating.
Here is how the additional input will look like:

As a result, the input will consist of pixels of each particular sample and a one-hot encoded label that will determine the class of this sample.

Note

You can find the source code via the following Link. If you want to run the code or even change some components, you can copy the notebook and work with the copy.

We can see that the generated images have good quality, and finally, we can control what we generate.
However, we have to take into account that the MNIST dataset is quite simple. If we want to generate more complex data, we have to consider more advanced architectures like GANs or their variations.