Update app.py

app.py

@@ -6,69 +6,78 @@ from huggingface_hub import hf_hub_download
 import os
 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):
-    super().__init__()
-    #encoder
-    self.img_2hid = nn.Linear(input_dim, h_dim)
-    self.hid_2mu = nn.Linear(h_dim, z_dim)
-    self.hid_2sigma = nn.Linear(h_dim, z_dim)
-
-    #decoder
-    self.z_2hid = nn.Linear(z_dim, h_dim)
-    self.hid_2img = nn.Linear(h_dim, input_dim)
-
-    self.relu = nn.ReLU()
-    #self.sigmoid = nn.sigmoid()
-
-  def encode(self, x):
-    h = self.relu(self.img_2hid(x))
-    mu, sigma = self.hid_2mu(h), self.hid_2sigma(h)
-    return mu, sigma
-
-  def decode(self, z):
-    h = self.relu(self.z_2hid(z))
-    return torch.sigmoid(self.hid_2img(h))
-
-  def forward(self, x):
-    mu, sigma = self.encode(x)
-    epsilon = torch.randn_like(sigma)
-    z_reparametrized = mu + sigma * epsilon
-    x_reconstructed = self.decode(z_reparametrized)
-    return x_reconstructed, mu, sigma
+    def __init__(self, input_dim, h_dim=400, z_dim=40):
+        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)
+        self.hid_2sigma = nn.Linear(h_dim, z_dim)
+
+        # Decoder
+        self.z_2hid = nn.Linear(z_dim, h_dim)
+        self.hid_2img = nn.Linear(h_dim, input_dim)
+
+        self.relu = nn.ReLU()
+
+    def encode(self, x):
+        h = self.relu(self.img_2hid(x))
+        mu = self.hid_2mu(h)
+        sigma = self.hid_2sigma(h)
+        return mu, sigma
+
+    def decode(self, z):
+        h = self.relu(self.z_2hid(z))
+        return torch.sigmoid(self.hid_2img(h))
+
+    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
 
 # -----------------------------
-# 1. CARGAR MODELO DESDE HUGGING FACE
+# 2. CARGAR MODELO DESDE HUGGING FACE
 # -----------------------------
-# Reemplaza estos datos con los tuyos
-REPO_ID = "Bmo411/VAE"  # <-- cámbialo por el tuyo
+REPO_ID = "Bmo411/VAE"  # ← tu repo
 MODEL_FILENAME = "vae_complete_model.pth"
 
-# Descargar modelo automáticamente
+# Descargar el modelo
 model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_FILENAME)
 
-# Inicializar modelo y cargar pesos
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = VAE(100*100)
+
+# Registrar la clase si se guardó como modelo completo
 torch.serialization.add_safe_globals({"VAE": VAE})
+
+# Cargar el modelo completo
 model = torch.load(model_path, map_location=device, weights_only=False)
 model.to(device)
 model.eval()
 
 # -----------------------------
-# 3. FUNCIÓN PARA GENERAR IMAGEN
+# 3. GENERAR IMAGEN ALEATORIA
 # -----------------------------
 def generate_image(z_dim=40):
     with torch.no_grad():
+        # Muestra del espacio latente
         z = torch.randn(1, z_dim).to(device)
-        out = model.decode(z)
-        out = torch.sigmoid(out)
+        out = model.decode(z)  # tamaño: (1, 10000)
+
+        # Convertir a forma imagen (1, 1, 100, 100)
         out = out.view(1, 1, 100, 100)
 
+        # Guardar imagen temporal
         output_path = "generated_sample.png"
         save_image(out, output_path)
 
-        img = Image.open(output_path)
+        # Leer imagen para mostrar en Gradio
+        img = Image.open(output_path).convert("L")
         return img
 
 # -----------------------------