Update app.py

app.py

@@ -5,6 +5,36 @@ from PIL import Image
 from huggingface_hub import hf_hub_download
 import os
 
+class VAE(nn.Module):
+  def __init__(self, input_dim, h_dim=400, z_dim=20):
+    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
 
 # -----------------------------
 # 1. CARGAR MODELO DESDE HUGGING FACE