Update app.py

app.py

@@ -1,34 +1,35 @@
-import torch
-import numpy as np
-import gradio as gr
-import matplotlib.pyplot as plt
-from model import Generator
-
-# Load the model
-device = torch.device("cpu")
-generator = Generator()
-generator.load_state_dict(torch.load("generator_digit.pth", map_location=device))
-generator.eval()
-
-def generate_images(digit):
-    noise = torch.randn(5, 100)
-    labels = torch.tensor([digit] * 5)
-    with torch.no_grad():
-        images = generator(noise, labels).squeeze().numpy()
-
-    # Create a grid of images (5 horizontal)
-    fig, axs = plt.subplots(1, 5, figsize=(10, 2))
-    for i in range(5):
-        axs[i].imshow(images[i], cmap='gray')
-        axs[i].axis('off')
-    return fig
-
-interface = gr.Interface(
-    fn=generate_images,
-    inputs=gr.inputs.Slider(0, 9, step=1, label="Digit (0–9)"),
-    outputs=gr.outputs.Image(type="plot"),
-    title="MNIST Digit Generator",
-    description="Generate 5 images of the digit using a conditional GAN trained on MNIST."
-)
-
-interface.launch()
+import torch
+import numpy as np
+import gradio as gr
+import matplotlib.pyplot as plt
+from model import Generator
+
+# Load the model
+device = torch.device("cpu")
+generator = Generator()
+generator.load_state_dict(torch.load("generator_digit.pth", map_location=device))
+generator.eval()
+
+def generate_images(digit):
+    noise = torch.randn(5, 100)
+    labels = torch.tensor([digit] * 5)
+    with torch.no_grad():
+        images = generator(noise, labels).squeeze().numpy()
+
+    # Plot the 5 images in one figure
+    fig, axs = plt.subplots(1, 5, figsize=(10, 2))
+    for i in range(5):
+        axs[i].imshow(images[i], cmap='gray')
+        axs[i].axis('off')
+    return fig
+
+# Gradio Interface using modern syntax
+demo = gr.Interface(
+    fn=generate_images,
+    inputs=gr.Slider(0, 9, step=1, label="Digit (0–9)"),
+    outputs=gr.Plot(label="Generated Images"),
+    title="MNIST Digit Generator",
+    description="Generates 5 handwritten images of the selected digit using a trained GAN."
+)
+
+demo.launch()