VAEを用いて画像をマッピングしたい。

Question

前提・実現したいこと

VAEを用いて自作のデータセットからマッピングを行おうとしています。
調べたところ、同じような事をしている記事があったので、それを参考にやったのですがうまくできませんでした。
参考サイト

発生している問題・エラーメッセージ

下記ソースコードを実行したところ、エラーは出なかったのですが、マッピングができておらず真っ暗な画像が表示されてしまいました。

該当のソースコード

データセット作成部分

import glob
import cv2
import matplotlib.pyplot
import numpy as np
x=glob.glob('{}/*.jpg'.format("/content/gdrive/My Drive/test/test1"))
image = []
for i in x:
  img = cv2.imread(i)
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  img = cv2.resize(img,(256,256))
  img = np.asarray(img)
  image.append(img)
  image.append(img[:, ::-1])

メイン

import tensorflow as tf
tf.compat.v1.disable_eager_execution()
import matplotlib.pyplot as plt
from keras import optimizers
%matplotlib inline
from keras.datasets import mnist
import keras
from keras import layers
from keras import backend as K
from keras.models import Model
import numpy as np
import glob
from keras.utils import np_utils
from PIL import Image
from scipy.stats import norm

image = np.asarray(image)
(x_train), (x_test) = image[:-32],image[-32:]
x_train = x_train.astype('float32') / 255.
x_train = np.reshape(x_train, (len(x_train), 256, 256, 3)) 
print(x_train.shape)
x_test = x_test.astype('float32') / 255.
x_test = np.reshape(x_test, (len(x_test), 256, 256, 3))

fig = plt.figure(figsize=(10, 10))
fig.subplots_adjust(left=0, right=1, bottom=0, top=0.5, hspace=0.01, wspace=0.01)
for i in range(100):
    ax = fig.add_subplot(10, 10, i + 1, xticks=[], yticks=[])
    ax.imshow(x_train[i].reshape((256, 256,3)), cmap='gray')


K.clear_session()

img_shape = (256, 256,3)
epochs = 1
batch_size = 32
latent_dim = 2

input_img = keras.Input(shape=img_shape)

x = layers.Conv2D(3, 3, padding='same',activation='relu')(input_img)
x = layers.Conv2D(32, 3,
                   padding='same',strides=2,activation='relu')(x)
x = layers.Conv2D(32, 3,padding='same',activation='relu')(x)
x = layers.Conv2D(32, 3, padding='same',activation='relu')(x)
shape_before_flattening = K.int_shape(x)

x = layers.Flatten()(x)
x = layers.Dense(128, activation='relu')(x)

z_mean = layers.Dense(latent_dim)(x)
z_log_var = layers.Dense(latent_dim)(x)
def sampling(args):
    z_mean, z_log_var = args
    epsilon = K.random_normal(shape=(K.shape(z_mean)[0], latent_dim),
                              mean=0., stddev=1.)
    return z_mean + K.exp(z_log_var) * epsilon

z = layers.Lambda(sampling)([z_mean, z_log_var])
encoder = Model(input_img,z)

x = layers.Dense(128, activation='relu')(decoder_input)
x = layers.Dense(shape_before_flattening[1]*shape_before_flattening[2]*shape_before_flattening[3], activation='relu')(x)

print(shape_before_flattening[1:])
x = layers.Reshape(shape_before_flattening[1:])(x)
x = layers.Conv2DTranspose(32, 3,padding='same',activation='relu',)(x)
x = layers.Conv2DTranspose(32, 3,padding='same',activation='relu',)(x)
x = layers.Conv2DTranspose(32, 3,padding='same',strides=2,activation='relu',)(x)
x = layers.Conv2D(3, 3,padding='same',activation='relu',)(x)
decoder = Model(decoder_input, x)

z_decoded = decoder(z)
class CustomVariationalLayer(keras.layers.Layer):

    def vae_loss(self, x, z_decoded):
        x = K.flatten(x)
        z_decoded = K.flatten(z_decoded)
        xent_loss = keras.metrics.binary_crossentropy(x, z_decoded)
        kl_loss = -5e-4 * K.mean(
            1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
        return K.mean(xent_loss + kl_loss)

    def call(self, inputs):
        x = inputs[0]
        z_decoded = inputs[1]
        loss = self.vae_loss(x, z_decoded)
        self.add_loss(loss, inputs=inputs)
        return x

y = CustomVariationalLayer()([input_img, z_decoded])
vae = Model(input_img, y)
sgd = optimizers.SGD(lr=0.0001, momentum=0.9)
adm = optimizers.Adam(lr=0.0001)
vae.compile(optimizer='rmsprop', loss=None)
vae.summary()
callbacks = [
   keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.5, patience=5, verbose=0, mode='auto'),
]
history = vae.fit(x=x_train, y=None,
        shuffle=True,
        epochs=epochs,
        batch_size=batch_size,
        callbacks=callbacks,
        validation_data=(x_test, None))
with open('loss.txt', 'a') as f:
    for loss in history.history['loss']:
        f.write(str(loss) + '\r')

n = 15
digit_size = 256
figure = np.zeros((digit_size*n, digit_size*n,3))
grid_x = norm.ppf(np.linspace(0.05, 0.95, n))
grid_y = norm.ppf(np.linspace(0.05, 0.95, n))

for i, yi in enumerate(grid_y):
    for j, xi in enumerate(grid_x):
        z_sample = np.array([[xi, yi]])
        z_sample = np.tile(z_sample, batch_size).reshape(batch_size, 2)
        x_decoded = decoder.predict(z_sample, batch_size)
        digit = x_decoded[0].reshape(digit_size,digit_size,3)
        figure[i*digit_size:(i+1)*digit_size,
               j*digit_size:(j+1)*digit_size] = digit
plt.figure(figsize=(10,10))
plt.imshow(figure, cmap='Greys_r')
plt.show()

試したこと

上記のソースコードを実行してtest1の中にある画像で二次元マッピングを行おうとしました。