Create app.py

app.py

@@ -0,0 +1,79 @@
+# app.py
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, transforms
+import gradio as gr
+import numpy as np
+from PIL import Image
+
+# 🧱 VAE architecture
+class VAE(nn.Module):
+    def __init__(self):
+        super(VAE, self).__init__()
+        self.fc1 = nn.Linear(784, 400)
+        self.fc21 = nn.Linear(400, 20)
+        self.fc22 = nn.Linear(400, 20)
+        self.fc3 = nn.Linear(20, 400)
+        self.fc4 = nn.Linear(400, 784)
+
+    def encode(self, x):
+        h1 = F.relu(self.fc1(x))
+        return self.fc21(h1), self.fc22(h1)
+
+    def reparam(self, mu, logvar):
+        std = torch.exp(0.5 * logvar)
+        eps = torch.randn_like(std)
+        return mu + eps * std
+
+    def decode(self, z):
+        return torch.sigmoid(self.fc4(F.relu(self.fc3(z))))
+
+    def forward(self, x):
+        mu, logvar = self.encode(x.view(-1, 784))
+        z = self.reparam(mu, logvar)
+        return self.decode(z), mu, logvar
+
+# 🧪 Loss function
+def vae_loss(recon_x, x, mu, logvar):
+    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784), reduction='sum')
+    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
+    return BCE + KLD
+
+# 📦 Load MNIST and train
+def train_vae():
+    vae = VAE()
+    optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)
+    loader = torch.utils.data.DataLoader(
+        datasets.MNIST('.', train=True, download=True,
+                       transform=transforms.ToTensor()),
+        batch_size=128, shuffle=True)
+    for epoch in range(1):
+        for x, _ in loader:
+            optimizer.zero_grad()
+            recon, mu, logvar = vae(x)
+            loss = vae_loss(recon, x, mu, logvar)
+            loss.backward()
+            optimizer.step()
+    return vae
+
+vae_model = train_vae()
+
+# 🎨 Gradio UI
+def generate(latent1=0.0, latent2=0.0):
+    z = torch.tensor([[latent1]*10 + [latent2]*10])
+    img = vae_model.decode(z).view(28, 28).detach().numpy()
+    img = (img * 255).astype(np.uint8)
+    return Image.fromarray(img, mode='L')
+
+demo = gr.Interface(
+    fn=generate,
+    inputs=[
+        gr.Slider(-3, 3, 0, label="Latent dim 1"),
+        gr.Slider(-3, 3, 0, label="Latent dim 2")
+    ],
+    outputs=gr.Image(type="pil", label="Generated Digit"),
+    title="🧠 VAE Digit Generator"
+)
+
+demo.launch()