Update app.py

app.py

@@ -1,65 +1,61 @@
-# gradio_app.py
-
-import torch
-import torch.nn as nn
-import gradio as gr
-import numpy as np
-from torchvision.utils import make_grid
-from torchvision.transforms import ToPILImage
-
-# === Match the Generator architecture ===
-
-class Generator(nn.Module):
-    def __init__(self, latent_dim, label_dim, img_shape):
-        super().__init__()
-        input_dim = latent_dim + label_dim
-        self.model = nn.Sequential(
-            nn.Linear(input_dim, 128),
-            nn.ReLU(True),
-            nn.Linear(128, 256),
-            nn.BatchNorm1d(256),
-            nn.ReLU(True),
-            nn.Linear(256, 512),
-            nn.BatchNorm1d(512),
-            nn.ReLU(True),
-            nn.Linear(512, int(torch.prod(torch.tensor(img_shape)))),
-            nn.Tanh()
-        )
-        self.img_shape = img_shape
-
-    def forward(self, noise, labels):
-        x = torch.cat((noise, labels), dim=1)
-        img = self.model(x)
-        return img.view(img.size(0), *self.img_shape)
-
-# === Load Model ===
-latent_dim = 100
-label_dim = 10
-img_shape = (1, 28, 28)
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-generator = Generator(latent_dim, label_dim, img_shape).to(device)
-generator.load_state_dict(torch.load("models/mnist_generator.pth", map_location=device))
-generator.eval()
-
-# === Generate Function ===
-def generate_images(digit):
-    digit = int(digit)
-    z = torch.randn(5, latent_dim, device=device)
-    labels = torch.eye(label_dim)[[digit]*5].to(device)
-    with torch.no_grad():
-        gen_imgs = generator(z, labels).cpu()
-    gen_imgs = (gen_imgs + 1) / 2  # de-normalize to [0,1]
-    images = [ToPILImage()(img.squeeze()) for img in gen_imgs]
-    return images
-
-# === Gradio App ===
-iface = gr.Interface(
-    fn=generate_images,
-    inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Choose a digit"),
-    outputs=[gr.Image(type="pil", label=f"Image {i+1}") for i in range(5)],
-    title="MNIST Digit Generator"
-)
-
-iface.launch()
-
+import torch
+import torch.nn as nn
+import gradio as gr
+import numpy as np
+from torchvision.transforms import ToPILImage
+
+# === Generator Model Definition ===
+class Generator(nn.Module):
+    def __init__(self, latent_dim, label_dim, img_shape):
+        super().__init__()
+        input_dim = latent_dim + label_dim
+        self.model = nn.Sequential(
+            nn.Linear(input_dim, 128),
+            nn.ReLU(True),
+            nn.Linear(128, 256),
+            nn.BatchNorm1d(256),
+            nn.ReLU(True),
+            nn.Linear(256, 512),
+            nn.BatchNorm1d(512),
+            nn.ReLU(True),
+            nn.Linear(512, int(torch.prod(torch.tensor(img_shape)))),
+            nn.Tanh()
+        )
+        self.img_shape = img_shape
+
+    def forward(self, noise, labels):
+        x = torch.cat((noise, labels), dim=1)
+        img = self.model(x)
+        return img.view(img.size(0), *self.img_shape)
+
+# === Load Model ===
+latent_dim = 100
+label_dim = 10
+img_shape = (1, 28, 28)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+generator = Generator(latent_dim, label_dim, img_shape).to(device)
+generator.load_state_dict(torch.load("models/mnist_generator.pth", map_location=device))
+generator.eval()
+
+# === Generate Function ===
+def generate_images(digit: str):
+    digit = int(digit)
+    z = torch.randn(5, latent_dim, device=device)
+    labels = torch.eye(label_dim, device=device)[[digit] * 5]
+    with torch.no_grad():
+        gen_imgs = generator(z, labels).cpu()
+    gen_imgs = (gen_imgs + 1) / 2  # Normalize to [0, 1]
+    images = [ToPILImage()(img.squeeze(0)) for img in gen_imgs]
+    return images
+
+# === Gradio Interface ===
+iface = gr.Interface(
+    fn=generate_images,
+    inputs=gr.Dropdown(choices=[str(i) for i in range(10)], label="Pick a digit"),
+    outputs=[gr.Image(type="pil") for _ in range(5)],
+    title="MNIST Digit Generator",
+    description="Select a digit from 0–9 to generate 5 synthetic handwritten digits using a GAN trained on MNIST."
+)
+
+iface.launch()