Update app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import gradio as gr
 from PIL import Image
 import os
+import math
 
 # Define your Generator architecture - with ngf=128 to match your training parameters
 class Generator(nn.Module):
@@ -64,32 +65,54 @@ except Exception as e:
 model.eval()
 print(f"Model initialized: {model is not None}")
 
-def generate_images(random_seed=42):
-    """Generate images using the DCGAN model"""
+def create_image_grid(images, rows, cols):
+    """Create a grid of images"""
+    w, h = images[0].size
+    grid = Image.new('RGB', size=(cols*w, rows*h))
+    
+    for i, image in enumerate(images):
+        grid.paste(image, box=(i%cols*w, i//cols*h))
+    return grid
+
+def generate_multiple_images(random_seed=42, num_images=4):
+    """Generate multiple images using the DCGAN model"""
     # Set seed for reproducibility
     torch.manual_seed(random_seed)
     
-    # Generate random noise
-    noise = torch.randn(1, 100, 1, 1, device=device)
+    # Generate multiple images
+    images = []
+    for i in range(num_images):
+        # Generate random noise with different seeds
+        noise = torch.randn(1, 100, 1, 1, device=device)
+        
+        # Generate fake image
+        with torch.no_grad():
+            fake_image = model(noise).detach().cpu()
+        
+        # Convert tensor to image
+        fake_img = fake_image * 0.5 + 0.5  # unnormalize
+        fake_img = fake_img.squeeze(0).permute(1, 2, 0).numpy()
+        fake_img = np.clip(fake_img * 255, 0, 255).astype(np.uint8)
+        images.append(Image.fromarray(fake_img))
     
-    # Generate fake image
-    with torch.no_grad():
-        fake_images = model(noise).detach().cpu()
+    # Create a grid of images
+    rows = int(math.sqrt(num_images))
+    cols = int(math.ceil(num_images / rows))
+    grid = create_image_grid(images, rows, cols)
     
-    # Convert tensor to image
-    fake_img = fake_images * 0.5 + 0.5  # unnormalize
-    fake_img = fake_img.squeeze(0).permute(1, 2, 0).numpy()
-    fake_img = np.clip(fake_img * 255, 0, 255).astype(np.uint8)
-    return Image.fromarray(fake_img)
+    return grid
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=generate_images,
-    inputs=gr.Slider(minimum=1, maximum=100, step=1, default=42, label="Random Seed"),
-    outputs=gr.Image(type="pil", label="Generated Computer Mouse"),
+    fn=generate_multiple_images,
+    inputs=[
+        gr.Slider(minimum=1, maximum=100, step=1, default=42, label="Random Seed"),
+        gr.Slider(minimum=1, maximum=16, step=1, default=4, label="Number of Images")
+    ],
+    outputs=gr.Image(type="pil", label="Generated Computer Mice"),
     title="DCGAN Computer Mouse Generator",
-    description="Generate unique computer mouse designs using a DCGAN model trained on computer mice images using ngf=128 and ndf=128.",
-    examples=[[42], [23], [7], [99]]
+    description="Generate multiple unique computer mouse designs using a DCGAN model.",
+    examples=[[42, 4], [23, 9], [7, 16]]
 )
 
 # Launch the app