VAE generates bad images. due to unbalanced loss functions?

I'm training a variational autoencoder on CelebA dataset using TensorFlow.keras

The problem I'm facing is that the generated images are not diverse enough and look kinda bad.

(new) Example:

What I think:

• it's bad because the reconstruction and KL loss are unbalanced.
• I read this question and followed its solution - read about KL annealing and tried to implement it myself but didn't work.

Note:

• It's my first time working with autoencoders so maybe I missed something obvious.

• it would be super appreciated if you could give a programmatic/technical solution and not a theoretical one with equations and complicated math

The loss function:

def r_loss(self, y_true, y_pred):
    return K.mean(K.square(y_true - y_pred), axis=[1, 2, 3])

def kl_loss(self, y_true, y_pred):
    return  -0.5 * K.sum(1 + self.sd_layer - K.square(self.mean_layer) - K.exp(self.sd_layer), axis=1)

def total_loss(self, y_true, y_pred):
    return K.mean(self.r_loss(y_true, y_pred) + self.kl_loss(y_true, y_pred))

The encoder:

    def build_encoder(self):

        conv_filters = [32, 64, 64, 64]
        conv_kernel_size = [3, 3, 3, 3]
        conv_strides = [2, 2, 2, 2]

        # Number of Conv layers
        n_layers = len(conv_filters)

        # Define model input
        x = self.encoder_input

        # Add convolutional layers
        for i in range(n_layers):
            x = Conv2D(filters=conv_filters[i],
                       kernel_size=conv_kernel_size[i],
                       strides=conv_strides[i],
                       padding='same',
                       name='encoder_conv_' + str(i)
                       )(x)
            if self.use_batch_norm: # True
                x = BatchNormalization()(x)

            x = LeakyReLU()(x)

            if self.use_dropout: # False
                x = Dropout(rate=0.25)(x)

        # Required for reshaping latent vector while building Decoder
        self.shape_before_flattening = K.int_shape(x)[1:]

        x = Flatten()(x)

        self.mean_layer = Dense(self.encoder_output_dim, name='mu')(x)
        self.sd_layer = Dense(self.encoder_output_dim, name='log_var')(x)


        # Defining a function for sampling
        def sampling(args):
            mean_mu, log_var = args
            epsilon = K.random_normal(shape=K.shape(mean_mu), mean=0., stddev=1.)
            return mean_mu + K.exp(log_var / 2) * epsilon

            # Using a Keras Lambda Layer to include the sampling function as a layer

        # in the model
        encoder_output = Lambda(sampling, name='encoder_output')([self.mean_layer, self.sd_layer])

        return Model(self.encoder_input, encoder_output, name=VAE_Encoder)

The decoder:

def build_decoder(self):
    conv_filters = [64, 64, 32, 3]
    conv_kernel_size = [3, 3, 3, 3]
    conv_strides = [2, 2, 2, 2]

    n_layers = len(conv_filters)

    # Define model input
    decoder_input = self.decoder_input

    # To get an exact mirror image of the encoder
    x = Dense(np.prod(self.shape_before_flattening))(decoder_input)
    x = Reshape(self.shape_before_flattening)(x)

    # Add convolutional layers
    for i in range(n_layers):
        x = Conv2DTranspose(filters=conv_filters[i],
                            kernel_size=conv_kernel_size[i],
                            strides=conv_strides[i],
                            padding='same',
                            name='decoder_conv_' + str(i)
                            )(x)

        # Adding a sigmoid layer at the end to restrict the outputs
        # between 0 and 1
        if i  n_layers - 1:
            x = LeakyReLU()(x)
        else:
            x = Activation('sigmoid')(x)

    # Define model output
    self.decoder_output = x

    return Model(decoder_input, self.decoder_output, name=VAE_Decoder)

The combined model:

def build_autoencoder(self):
    self.encoder = self.build_encoder()
    self.decoder = self.build_decoder()

    # Input to the combined model will be the input to the encoder.
    # Output of the combined model will be the output of the decoder.
    self.autoencoder = Model(self.encoder_input, self.decoder(self.encoder(self.encoder_input)),
                             name=Variational_Auto_Encoder)

    self.autoencoder.compile(optimizer=self.adam_optimizer, loss=self.total_loss,
                             metrics=[self.total_loss],
                             experimental_run_tf_function=False)
    self.autoencoder.summary()

EDIT:

the latent size is 256 and the sample method is as follows;

def generate(self, image=None):
    if not os.path.exists(self.sample_dir):
        os.makedirs(self.sample_dir)
    if image is None:
        img = np.random.normal(size=(9, self.encoder_output_dim))

        prediction = self.decoder.predict(img)

        op = np.vstack((np.hstack((prediction[0], prediction[1], prediction[2])),
                        np.hstack((prediction[3], prediction[4], prediction[5])),
                        np.hstack((prediction[6], prediction[7], prediction[8]))))
        print(op.shape)
        op = cv2.resize(op, (self.input_size * 9, self.input_size * 9), interpolation=cv2.INTER_AREA)
        op = cv2.cvtColor(op, cv2.COLOR_BGR2RGB)
        cv2.imshow(generated, op)
        cv2.imwrite(self.sample_dir + generated + str(r(0, 9999)) + .jpg, (op * 255).astype(uint8))

    else:
        img = cv2.imread(image, cv2.IMREAD_UNCHANGED)
        img = cv2.resize(img, (self.input_size, self.input_size), interpolation=cv2.INTER_AREA)
        img = img.astype(float32)
        img = img / 255

        prediction = self.autoencoder.predict(img.reshape(1, self.input_size, self.input_size, 3))
        img = cv2.resize(prediction[0][:, :, ::-1], (960, 960), interpolation=cv2.INTER_AREA)

        cv2.imshow(prediction, img)

        cv2.imwrite(self.sample_dir + generated + str(r(0, 9999)) + .jpg, (img * 255).astype(uint8))