ulepsz mój trening modelu: import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Reshape, Conv2DTranspose, Conv2D, Flatten
from tensorflow.keras.models import Model
from tensorflow.keras.datasets import cifar10
import numpy as np
import matplotlib.pyplot as plt
from tensorflow.keras.layers import Layer

# Sampling Layer
class Sampling(Layer):
    def call(self, inputs):
        z_mean, z_log_var = inputs
        epsilon = tf.keras.backend.random_normal(shape=tf.shape(z_mean))
        return z_mean + tf.exp(0.5 * z_log_var) * epsilon

# Encoder
def build_conv_encoder(input_shape, latent_dim):
    input_img = Input(shape=input_shape)
    x = Conv2D(32, kernel_size=3, strides=2, activation='relu', padding='same')(input_img)
    x = Conv2D(64, kernel_size=3, strides=2, activation='relu', padding='same')(x)
    x_shape = tf.keras.backend.int_shape(x)[1:]  # Shape before flattening
    x = Flatten()(x)
    x = Dense(128, activation='relu')(x)
    z_mean = Dense(latent_dim, name='z_mean')(x)
    z_log_var = Dense(latent_dim, name='z_log_var')(x)
    z = Sampling()([z_mean, z_log_var])
    encoder = Model(input_img, [z_mean, z_log_var, z], name='encoder')
    return encoder, x_shape

# Decoder
def build_conv_decoder(x_shape, latent_dim, output_shape):
    latent_inputs = Input(shape=(latent_dim,))
    x = Dense(np.prod(x_shape), activation='relu')(latent_inputs)
    x = Reshape(x_shape)(x)
    x = Conv2DTranspose(64, kernel_size=3, strides=2, activation='relu', padding='same')(x)
    x = Conv2DTranspose(32, kernel_size=3, strides=2, activation='relu', padding='same')(x)
    outputs = Conv2DTranspose(output_shape[2], kernel_size=3, activation='sigmoid', padding='same')(x)
    decoder = Model(latent_inputs, outputs, name='decoder')
    return decoder

# VAE Model
class VAE(Model):
    def __init__(self, encoder, decoder, kl_weight=1.0, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.kl_weight = kl_weight
        self.reconstruction_loss_tracker = tf.keras.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = tf.keras.metrics.Mean(name="kl_loss")

    def call(self, inputs):
        z_mean, z_log_var, z = self.encoder(inputs)
        reconstruction = self.decoder(z)
        return reconstruction

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(
                tf.reduce_sum(
                    tf.keras.losses.binary_crossentropy(data, reconstruction),
                    axis=(1, 2)
                )
            )
            kl_loss = -0.5 * tf.reduce_mean(
                tf.reduce_sum(
                    z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1,
                    axis=1
                )
            )
            total_loss = reconstruction_loss + self.kl_weight * kl_loss

        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))

        # Update metrics
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)

        return {"loss": total_loss,
                "reconstruction_loss": self.reconstruction_loss_tracker.result(),
                "kl_loss": self.kl_loss_tracker.result()}

    def test_step(self, data):
        z_mean, z_log_var, z = self.encoder(data)
        reconstruction = self.decoder(z)
        reconstruction_loss = tf.reduce_mean(
            tf.reduce_sum(
                tf.keras.losses.binary_crossentropy(data, reconstruction),
                axis=(1, 2)
            )
        )
        kl_loss = -0.5 * tf.reduce_mean(
            tf.reduce_sum(
                z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1,
                axis=1
            )
        )
        total_loss = reconstruction_loss + self.kl_weight * kl_loss

        # Update metrics
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)

        return {"loss": total_loss,
                "reconstruction_loss": self.reconstruction_loss_tracker.result(),
                "kl_loss": self.kl_loss_tracker.result()}

    @property
    def metrics(self):
        return [self.reconstruction_loss_tracker, self.kl_loss_tracker]

# Train VAE
def train_vae(input_shape, latent_dim, x_train, x_val, epochs=30, batch_size=128, dataset_name='Dataset'):
    encoder, x_shape = build_conv_encoder(input_shape, latent_dim)
    decoder = build_conv_decoder(x_shape, latent_dim, input_shape)
    vae = VAE(encoder, decoder, kl_weight=0.1)
    vae.compile(optimizer='adam', loss='mse')
    history = vae.fit(
        x_train,
        epochs=epochs,
        batch_size=batch_size,
        validation_data=(x_val, None)
    )

    # Plot training losses
    plt.figure(figsize=(8, 6))
    plt.plot(history.history['loss'], label='Training Total Loss')
    plt.plot(history.history['reconstruction_loss'], label='Training Reconstruction Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.title(f'Training Loss on {dataset_name}')
    plt.show()

    # Reconstruct test images
    x_reconstructed = vae.predict(x_val[:10])
    plot_reconstructions(x_val[:10], x_reconstructed, input_shape)
    return vae

# Plot original and reconstructed images
def plot_reconstructions(x_original, x_reconstructed, input_shape, num_images=10):
    plt.figure(figsize=(20, 4))
    for i in range(num_images):
        ax = plt.subplot(2, num_images, i + 1)
        plt.imshow(x_original[i])
        ax.axis('off')
        ax = plt.subplot(2, num_images, i + 1 + num_images)
        plt.imshow(x_reconstructed[i])
        ax.axis('off')
    plt.tight_layout()
    plt.show()

# Load CIFAR-10
(x_train_cifar, _), (x_test_cifar, _) = cifar10.load_data()
x_train_cifar = x_train_cifar.astype('float32') / 255.
x_test_cifar = x_test_cifar.astype('float32') / 255.

# Train VAE on CIFAR-10
vae_cifar = train_vae(
    input_shape=(32, 32, 3),
    latent_dim=32,
    x_train=x_train_cifar,
    x_val=x_test_cifar,
    epochs=30,
    batch_size=256,
    dataset_name='CIFAR-10'
)