Update app.py

app.py

@@ -10,9 +10,8 @@ import torch.nn as nn
 # 1. DEFINICIÓN DEL MODELO VAE
 # -----------------------------
 class VAE(nn.Module):
-    def __init__(self, input_dim, h_dim=400, z_dim=40):
+    def __init__(self, input_dim, h_dim=400, z_dim=20):  # NOTA: z_dim por defecto en 20
         super().__init__()
-        self.z_dim = z_dim
         # Encoder
         self.img_2hid = nn.Linear(input_dim, h_dim)
         self.hid_2mu = nn.Linear(h_dim, z_dim)
@@ -26,8 +25,7 @@ class VAE(nn.Module):
 
     def encode(self, x):
         h = self.relu(self.img_2hid(x))
-        mu = self.hid_2mu(h)
-        sigma = self.hid_2sigma(h)
+        mu, sigma = self.hid_2mu(h), self.hid_2sigma(h)
         return mu, sigma
 
     def decode(self, z):
@@ -37,47 +35,48 @@ class VAE(nn.Module):
     def forward(self, x):
         mu, sigma = self.encode(x)
         epsilon = torch.randn_like(sigma)
-        z_reparam = mu + sigma * epsilon
-        x_recon = self.decode(z_reparam)
-        return x_recon, mu, sigma
+        z_reparametrized = mu + sigma * epsilon
+        x_reconstructed = self.decode(z_reparametrized)
+        return x_reconstructed, mu, sigma
 
 # -----------------------------
-# 2. CARGAR MODELO DESDE HUGGING FACE
+# 2. CARGAR EL MODELO DESDE HUGGING FACE
 # -----------------------------
-REPO_ID = "Bmo411/VAE"  # ← tu repo
+REPO_ID = "Bmo411/VAE"  # <-- reemplaza con tu repo si cambia
 MODEL_FILENAME = "vae_complete_model.pth"
 
-# Descargar el modelo
+# Descargar modelo automáticamente
 model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Inicializar arquitectura del modelo
+input_dim = 100 * 100
+dummy_model = VAE(input_dim=input_dim, z_dim=20)  # la arquitectura base es necesaria para cargar pesos
 
-# Registrar la clase si se guardó como modelo completo
+# Permitir deserialización segura
 torch.serialization.add_safe_globals({"VAE": VAE})
 
-# Cargar el modelo completo
+# Cargar modelo completo (no solo pesos)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.to(device)
 model.eval()
 
+# Detectar z_dim automáticamente desde el decoder
+z_dim = model.z_2hid.in_features
+
 # -----------------------------
-# 3. GENERAR IMAGEN ALEATORIA
+# 3. FUNCIÓN PARA GENERAR IMAGEN
 # -----------------------------
-def generate_image(z_dim=40):
+def generate_image():
     with torch.no_grad():
-        # Muestra del espacio latente
         z = torch.randn(1, z_dim).to(device)
-        out = model.decode(z)  # tamaño: (1, 10000)
-
-        # Convertir a forma imagen (1, 1, 100, 100)
+        out = model.decode(z)
         out = out.view(1, 1, 100, 100)
 
-        # Guardar imagen temporal
         output_path = "generated_sample.png"
         save_image(out, output_path)
 
-        # Leer imagen para mostrar en Gradio
-        img = Image.open(output_path).convert("L")
+        img = Image.open(output_path).convert("L")  # Convertir a escala de grises
         return img
 
 # -----------------------------
@@ -85,10 +84,10 @@ def generate_image(z_dim=40):
 # -----------------------------
 iface = gr.Interface(
     fn=generate_image,
-    inputs=gr.Slider(10, 100, value=40, step=1, label="Dimensión latente (z_dim)"),
+    inputs=[],
     outputs="image",
     title="Generador de Imagen con VAE",
-    description="Genera una imagen aleatoria a partir del espacio latente del VAE entrenado."
+    description=f"Genera una imagen aleatoria desde el VAE entrenado. Dimensión latente del modelo detectada: {z_dim}"
 )
 
 iface.launch()