app.py

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

# ---------------------------------------------------------------------------
# 1. Model Architecture (Must match training)
# ---------------------------------------------------------------------------
class Generator(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 512),
            nn.ReLU(),
            nn.Linear(512, 1024),
            nn.ReLU(),
            nn.Linear(1024, 3 * 64 * 64),
            nn.Tanh()
        )

    def forward(self, z):
        img = self.model(z)
        return img.view(-1, 3, 64, 64)

# ---------------------------------------------------------------------------
# 2. Configuration & Model Loading
# ---------------------------------------------------------------------------
LATENT_DIM = 100
MODEL_PATH = "generator_epoch_40.pth"
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def load_generator():
    model = Generator()
    if os.path.exists(MODEL_PATH):
        try:
            state_dict = torch.load(MODEL_PATH, map_location=DEVICE)
            model.load_state_dict(state_dict)
            model.to(DEVICE)
            model.eval()
            print(f"Model loaded successfully from {MODEL_PATH}")
            return model
        except Exception as e:
            print(f"Error loading weights: {e}")
    else:
        print(f"Warning: {MODEL_PATH} not found. Running with random weights.")
    
    model.to(DEVICE)
    model.eval()
    return model

generator = load_generator()

# ---------------------------------------------------------------------------
# 3. Inference Function
# ---------------------------------------------------------------------------
def generate_images(num_images):
    with torch.no_grad():
        noise = torch.randn(int(num_images), LATENT_DIM).to(DEVICE)
        fake_images = generator(noise).detach().cpu()
        
        # Scale to [0, 1] for visualization
        grid = vutils.make_grid(fake_images, padding=2, normalize=True, nrow=4)
        
        # Convert to PIL Image
        # grid is (3, H, W)
        ndarr = grid.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy()
        return Image.fromarray(ndarr)

# ---------------------------------------------------------------------------
# 4. Gradio Interface
# ---------------------------------------------------------------------------
custom_css = """
.gradio-container {
    max-width: 800px !important;
    margin: 0 auto;
    background-color: #0f172a;
    color: white;
}
.main-header {
    text-align: center;
    padding: 2rem 0;
    background: linear-gradient(135deg, #334155 0%, #1e293b 100%);
    border-radius: 12px;
    margin-bottom: 2rem;
}
.gr-button-primary {
    background: linear-gradient(90deg, #f59e0b 0%, #d97706 100%) !important;
    border: none !important;
}
"""

with gr.Blocks(css=custom_css, title="Vanilla GAN Anime Face Generator") as demo:
    with gr.Column(elem_classes="main-header"):
        gr.Markdown(
            """
            # 🎨 Vanilla GAN: Anime Face Generator
            ### Generating unique anime characters using a Multi-Layer Perceptron GAN.
            """
        )
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### ⚙️ Generation Settings")
            num_slider = gr.Slider(
                minimum=1,
                maximum=16,
                step=1,
                value=4,
                label="Number of Faces",
                info="Choose how many faces to generate in a grid."
            )
            generate_btn = gr.Button("✨ Generate Faces", variant="primary")
            
            with gr.Accordion("About the Model", open=False):
                gr.Markdown(
                    """
                    - **Architecture**: Simple MLP Generator and Discriminator.
                    - **Training**: 60 epochs on Anime Face Dataset.
                    - **Input**: 100D Latent Noise.
                    - **Output**: 64x64 RGB Images.
                    """
                )
        
        with gr.Column():
            output_image = gr.Image(label="Generated Result", type="pil")

    gr.HTML(
        """
        <p style="text-align: center; color: #94a3b8; margin-top: 2rem; font-size: 0.8rem;">
        Created with PyTorch & Gradio | Hugging Face Spaces
        </p>
        """
    )

    generate_btn.click(
        fn=generate_images,
        inputs=[num_slider],
        outputs=[output_image]
    )

if __name__ == "__main__":
    demo.launch()