import gradio as gr
import tensorflow as tf
import numpy as np
from PIL import Image
from huggingface_hub import hf_hub_download

# Cargar el modelo DAE desde Hugging Face
modelo_path = hf_hub_download(repo_id="XTEP63/VAE", filename="VAE.keras")
vae_model = tf.keras.models.load_model(modelo_path)

# Función para generar una imagen con el VAE
def generate_image():
    latent_dim = 128  
    z_sample = np.random.normal(size=(1, latent_dim))  # Muestra aleatoria del espacio latente
    output = vae_model.predict(z_sample)  # Generar imagen
    output = np.squeeze(output, axis=0)  
    output = (output * 255).astype(np.uint8)  
    return Image.fromarray(output)

# Interfaz con Gradio
iface_vae = gr.Interface(
    fn=generate_image,
    inputs=None,
    outputs=gr.Image(type="pil"),
    title="Variational Autoencoder",
    description="Genera una imagen nueva basada en el espacio latente del VAE."
)

iface_vae.launch()