app.py

import gradio as gr
import torch
import torch.nn as nn
import numpy as np
from PIL import Image
import os

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Model parameters (must match your training)
nz = 100
ngf = 64
num_classes = 10

# Generator class (same as your training script)
class Generator(nn.Module):
    def __init__(self):
        super(Generator, self).__init__()
        
        self.label_embedding = nn.Embedding(num_classes, num_classes)
        
        self.main = nn.Sequential(
            nn.ConvTranspose2d(nz + num_classes, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            
            nn.ConvTranspose2d(ngf * 2, 1, 4, 2, 1, bias=False),
            nn.Tanh()
        )
        
        self.resize = nn.AdaptiveAvgPool2d((28, 28))
    
    def forward(self, noise, labels):
        label_embedding = self.label_embedding(labels)
        label_embedding = label_embedding.view(label_embedding.size(0), num_classes, 1, 1)
        input_tensor = torch.cat([noise, label_embedding], dim=1)
        output = self.main(input_tensor)
        output = self.resize(output)
        return output

# Load the trained model
@st.cache_resource
def load_model():
    generator = Generator().to(device)
    
    # Load the saved model
    if os.path.exists('mnist_gan_model.pth'):
        checkpoint = torch.load('mnist_gan_model.pth', map_location=device)
        generator.load_state_dict(checkpoint['generator_state_dict'])
        generator.eval()
        print("Model loaded successfully!")
    else:
        print("Warning: Model file not found!")
    
    return generator

# Initialize generator
generator = load_model()

# Generation function
def generate_digit_images(digit):
    """Generate 5 images of the specified digit"""
    
    digit = int(digit)
    num_images = 5
    
    with torch.no_grad():
        # Generate random noise
        noise = torch.randn(num_images, nz, 1, 1).to(device)
        labels = torch.full((num_images,), digit, dtype=torch.long).to(device)
        
        # Generate images
        generated_images = generator(noise, labels)
        
        # Convert to numpy and denormalize
        images = generated_images.cpu().numpy()
        images = (images + 1) / 2.0  # Denormalize from [-1, 1] to [0, 1]
        images = np.squeeze(images)  # Remove channel dimension
        
        # Convert to PIL Images for Gradio
        pil_images = []
        for img in images:
            # Convert to 0-255 range and uint8
            img_uint8 = (img * 255).astype(np.uint8)
            pil_img = Image.fromarray(img_uint8, mode='L')
            # Resize for better visibility
            pil_img = pil_img.resize((112, 112), Image.NEAREST)  # 4x upscale
            pil_images.append(pil_img)
        
        return pil_images

# Create Gradio interface
def create_app():
    with gr.Blocks(
        title="Handwritten Digit Generator", 
        theme=gr.themes.Soft(),
        css=".gradio-container {max-width: 700px; margin: auto;}"
    ) as app:
        
        gr.Markdown("# 🔢 Handwritten Digit Generator")
        gr.Markdown("Generate synthetic MNIST-like digit images using a trained GAN model. Select a digit (0-9) to generate 5 unique images.")
        
        with gr.Row():
            with gr.Column(scale=1):
                digit_input = gr.Dropdown(
                    choices=list(range(10)), 
                    value=2, 
                    label="Choose a digit to generate (0-9)",
                    interactive=True
                )
                generate_btn = gr.Button(
                    "🎨 Generate Images", 
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=2):
                gr.Markdown("### Generated Images")
                
                # Gallery to display 5 images
                image_gallery = gr.Gallery(
                    label="Generated Digit Images", 
                    show_label=False,
                    columns=5,
                    rows=1,
                    height=200,
                    object_fit="contain"
                )
        
        # Example section
        gr.Markdown("---")
        gr.Markdown("### How it works")
        gr.Markdown("""
        1. **Select** a digit from the dropdown (0-9)
        2. **Click** 'Generate Images' button  
        3. **View** 5 unique generated images of your chosen digit
        4. Each generation produces different variations of the same digit
        """)
        
        # Connect button to generation function
        generate_btn.click(
            fn=generate_digit_images,
            inputs=[digit_input],
            outputs=[image_gallery]
        )
        
        # Auto-generate on page load
        app.load(
            fn=generate_digit_images,
            inputs=[gr.Number(value=2, visible=False)],
            outputs=[image_gallery]
        )
        
        # Footer
        gr.Markdown("---")
        gr.Markdown("**🤖 Model**: Conditional GAN trained on MNIST | **⚡ Framework**: PyTorch + Gradio")
    
    return app

# Launch the app
if __name__ == "__main__":
    app = create_app()
    app.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False  # Set to False for Hugging Face deployment
    )