Update app.py

app.py

@@ -1,21 +1,24 @@
 import torch
 import numpy as np
-import matplotlib.pyplot as plt
 import gradio as gr
+import matplotlib.pyplot as plt
 
-# Define the same VAE architecture used during training
-class VAE(torch.nn.Module):
-    def __init__(self):
+# Conditional VAE definition (same as training)
+class CVAE(torch.nn.Module):
+    def __init__(self, latent_dim=20):
         super().__init__()
+        self.latent_dim = latent_dim
+        self.label_embed = torch.nn.Embedding(10, 10)
+
         self.encoder = torch.nn.Sequential(
-            torch.nn.Flatten(),
-            torch.nn.Linear(28*28, 400),
+            torch.nn.Linear(28*28 + 10, 400),
             torch.nn.ReLU(),
         )
-        self.mu = torch.nn.Linear(400, 20)
-        self.logvar = torch.nn.Linear(400, 20)
+        self.fc_mu = torch.nn.Linear(400, latent_dim)
+        self.fc_logvar = torch.nn.Linear(400, latent_dim)
+
         self.decoder = torch.nn.Sequential(
-            torch.nn.Linear(20, 400),
+            torch.nn.Linear(latent_dim + 10, 400),
             torch.nn.ReLU(),
             torch.nn.Linear(400, 28*28),
             torch.nn.Sigmoid()
@@ -26,34 +29,34 @@ class VAE(torch.nn.Module):
         eps = torch.randn_like(std)
         return mu + eps * std
 
-    def forward(self, x):
-        h = self.encoder(x)
-        mu, logvar = self.mu(h), self.logvar(h)
-        z = self.reparameterize(mu, logvar)
-        return self.decoder(z)
+    def decode(self, z, y):
+        y_embed = self.label_embed(y)
+        inputs = torch.cat([z, y_embed], dim=1)
+        return self.decoder(inputs)
 
-# Load model
-model = VAE()
+model = CVAE()
 model.load_state_dict(torch.load("cvae_mnist.pth", map_location='cpu'))
 model.eval()
 
-# Generation function for Gradio
-def generate_images(digit):
-    # For VAE, we ignore the digit and generate random samples
+# Image generation function
+def generate_digit_images(digit):
     images = []
     for _ in range(5):
         z = torch.randn(1, 20)
-        img = model.decoder(z).detach().numpy().reshape(28, 28)
+        y = torch.tensor([int(digit)])
+        with torch.no_grad():
+            out = model.decode(z, y)
+        img = out.view(28, 28).numpy()
         images.append((img * 255).astype(np.uint8))
     return images
 
-# Gradio interface
+# Launch Gradio app
 iface = gr.Interface(
-    fn=generate_images,
-    inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Choose a digit (ignored for now)"),
+    fn=generate_digit_images,
+    inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Choose a digit (0–9)"),
     outputs=[gr.Image(image_mode='L') for _ in range(5)],
-    title="Handwritten Digit Generator",
-    description="Select a digit (0–9) and generate 5 handwritten-style digits using a VAE trained on MNIST."
+    title="Conditional VAE Handwritten Digit Generator",
+    description="Generates 5 images of the digit you select (0–9) using a Conditional Variational Autoencoder trained on MNIST."
 )
 
-iface.launch()
+iface.launch()