app.py

import streamlit as st
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torch.utils.data import Dataset, DataLoader
import numpy as np
from PIL import Image
import io
import matplotlib.pyplot as pltF
import gc
import os
# Add these imports at the top of your file
from tinydb import TinyDB, Query
import base64
import datetime
import matplotlib.pyplot as plt


# Set page config with colorful theme
st.set_page_config(
    page_title="AI Art Generator - GAN Training Studio", 
    layout="wide",
    initial_sidebar_state="expanded",
    page_icon="🎨"
)

# --- Background Image Encoding ---
def get_base64_of_bin_file(bin_file):
    with open(bin_file, 'rb') as f:
        data = f.read()
    return base64.b64encode(data).decode()

def set_png_as_page_bg(bin_file):
    bin_str = get_base64_of_bin_file(bin_file)
    page_bg_img = f'''
    <style>
    .stApp {{
        background-image: url("data:image/png;base64,{bin_str}");
        background-size: cover;
        background-repeat: no-repeat;
        background-attachment: fixed;
        background-position: center;
        image-rendering: pixelated;
    }}
    </style>
    '''
    st.markdown(page_bg_img, unsafe_allow_html=True)

# Set the local background image
if os.path.exists("image.png"):
    set_png_as_page_bg("image.png")

# Custom CSS for magical, dreamy, transparent pixel image
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Press+Start+2P&display=swap');

    /* Make the main containers and header totally transparent */
    [data-testid="stAppViewContainer"], 
    [data-testid="stHeader"] {
        background-color: transparent !important;
    }

    /* Full Transparency sidebar */
    [data-testid="stSidebar"] {
        background-color: transparent !important;
        border-right: 1px solid rgba(255, 255, 255, 0.1);
    }
    
    /* Sidebar inner content background reset */
    [data-testid="stSidebar"] > div:first-child {
        background-color: transparent !important;
    }

    /* Main block container - Full Transparency */
    .block-container {
        background-color: transparent !important;
        border-radius: 20px;
        padding: 2.5rem 3.5rem !important;
        border: 1px solid rgba(255, 255, 255, 0.1);
        box-shadow: none !important;
    }
    
    /* Removed all text shadows entirely to avoid any dark/grey shading as requested */
    p, span, label, div, li, h2, h3, h4, h5, h6 {
        text-shadow: none !important;
    }

    .big-title {
        font-family: 'Press Start 2P', cursive !important;
        font-size: 2.5rem !important;
        /* Exact color match from pixel bg: Pink, Blue, Gold trail */
        background: linear-gradient(to right, #ff44cc, #44ccff, #ffcc33);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        text-align: center;
        margin-bottom: 2rem !important;
        text-shadow: none !important; 
        filter: none !important; 
    }

    /* Clean subtle subtitle */
    .main-subtitle h3 {
        color: #f0f0f0 !important;
        text-shadow: none !important;
        font-weight: 300 !important;
    }

    /* Fully Transparent inputs */
    div[data-baseweb="select"] > div, 
    input, 
    div[data-testid="stFileUploader"] {
        background-color: rgba(0, 0, 0, 0.2) !important; /* Extremely light touch */
        border: 1px solid rgba(255, 255, 255, 0.2) !important;
        border-radius: 8px;
    }
    
    /* Tab transparency */
    button[data-baseweb="tab"] {
        background-color: transparent !important;
        border: 1px solid rgba(255, 255, 255, 0.1) !important;
        border-radius: 10px 10px 0 0;
        margin-right: 5px;
    }
    
    button[data-baseweb="tab"][aria-selected="true"] {
        background-color: rgba(255,255,255,0.15) !important;
        border-bottom: 2px solid #4ecdc4 !important;
    }
    
    /* Make specific nested blocks more transparent to not stack too darkly */
    [data-testid="stVerticalBlock"] {
        background-color: transparent !important;
    }
</style>
""", unsafe_allow_html=True)




# Device configuration for memory efficiency
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class Generator128(nn.Module):
    """Generator capable of producing 128x128 images."""
    def __init__(self, latent_dim=100, img_channels=3, feature_maps=64):
        super().__init__()
        self.main = nn.Sequential(
            # Input: Z (latent vector)
            nn.ConvTranspose2d(latent_dim, feature_maps * 16, 4, 1, 0, bias=False),
            nn.BatchNorm2d(feature_maps * 16),
            nn.ReLU(True),
            # State size: (feature_maps*16) x 4 x 4
            nn.ConvTranspose2d(feature_maps * 16, feature_maps * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps * 8),
            nn.ReLU(True),
            # State size: (feature_maps*8) x 8 x 8
            nn.ConvTranspose2d(feature_maps * 8, feature_maps * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps * 4),
            nn.ReLU(True),
            # State size: (feature_maps*4) x 16 x 16
            nn.ConvTranspose2d(feature_maps * 4, feature_maps * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps * 2),
            nn.ReLU(True),
            # State size: (feature_maps*2) x 32 x 32
            nn.ConvTranspose2d(feature_maps * 2, feature_maps, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps),
            nn.ReLU(True),
            # State size: (feature_maps) x 64 x 64
            # --- NEW LAYER FOR 128x128 ---
            nn.ConvTranspose2d(feature_maps, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
            # Final state size: (img_channels) x 128 x 128
        )

    def forward(self, z):
        z = z.view(z.shape[0], -1, 1, 1)
        return self.main(z)




class Discriminator128(nn.Module):
    """Discriminator capable of handling 128x128 images."""
    def __init__(self, img_channels=3, feature_maps=64):
        super().__init__()
        self.main = nn.Sequential(
            # --- NEW LAYER FOR 128x128 INPUT ---
            # Input size: (img_channels) x 128 x 128
            nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps) x 64 x 64
            nn.Conv2d(feature_maps, feature_maps * 2, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 2, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*2) x 32 x 32
            nn.Conv2d(feature_maps * 2, feature_maps * 4, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 4, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*4) x 16 x 16
            nn.Conv2d(feature_maps * 4, feature_maps * 8, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 8, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*8) x 8 x 8
            nn.Conv2d(feature_maps * 8, feature_maps * 16, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 16, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*16) x 4 x 4
            nn.Conv2d(feature_maps * 16, 1, 4, 1, 0, bias=False),
        )

    def forward(self, img):
        return self.main(img)



# ===================== SIMPLE GAN MODELS =====================
class SimpleGenerator(nn.Module):
    """Lightweight generator network"""
    def __init__(self, latent_dim=100, img_channels=3, img_size=64):
        super().__init__()
        self.img_size = img_size
        self.img_channels = img_channels
        
        # Calculate initial size for reshape
        init_size = img_size // 4
        self.init_size = init_size
        
        self.linear = nn.Sequential(
            nn.Linear(latent_dim, 128 * init_size ** 2),
            nn.LeakyReLU(0.2)
        )
        
        self.conv_blocks = nn.Sequential(
            nn.BatchNorm2d(128),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 128, 3, stride=1, padding=1),
            nn.BatchNorm2d(128),
            nn.LeakyReLU(0.2),
            nn.Upsample(scale_factor=2),
            nn.Conv2d(128, 64, 3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.LeakyReLU(0.2),
            nn.Conv2d(64, img_channels, 3, stride=1, padding=1),
            nn.Tanh()
        )
    
    def forward(self, z):
        out = self.linear(z)
        out = out.view(out.shape[0], 128, self.init_size, self.init_size)
        img = self.conv_blocks(out)
        return img

class SimpleDiscriminator(nn.Module):
    """Lightweight discriminator network"""
    def __init__(self, img_channels=3, img_size=64):
        super().__init__()
        
        def discriminator_block(in_filters, out_filters, stride=2):
            return [
                nn.Conv2d(in_filters, out_filters, 3, stride, 1),
                nn.LeakyReLU(0.2),
                nn.Dropout2d(0.25)
            ]
        
        self.model = nn.Sequential(
            *discriminator_block(img_channels, 16, 2),
            *discriminator_block(16, 32, 2),
            *discriminator_block(32, 64, 2),
            *discriminator_block(64, 128, 2),
        )
        
        # Calculate the size after convolutions
        ds_size = img_size // 2 ** 4
        self.adv_layer = nn.Sequential(
            nn.Linear(128 * ds_size ** 2, 1),
            nn.Sigmoid()
        )
    
    def forward(self, img):
        out = self.model(img)
        out = out.view(out.shape[0], -1)
        validity = self.adv_layer(out)
        return validity

# ===================== DCGAN MODELS =====================
# ===================== DCGAN MODELS (DYNAMIC & IMPROVED) =====================
# ===================== DCGAN MODELS (DYNAMIC & IMPROVED) =====================
# ===================== DCGAN MODELS (DYNAMIC & IMPROVED) =====================
class Generator(nn.Module):
    """Powerful and FLEXIBLE DCGAN-style generator"""
    def __init__(self, latent_dim=100, img_channels=3, img_size=64, feature_maps=64):
        super().__init__()
        
        # Calculate the number of upsampling blocks needed to go from 4x4 to img_size
        num_blocks = int(np.log2(img_size) - 2)
        if num_blocks < 1:
            raise ValueError(f"Image size {img_size} is too small for this DCGAN architecture. Minimum is 8.")

        layers = []
        
        # Initial layer: from latent_dim to a 4x4 feature map
        out_features = feature_maps * (2 ** (num_blocks - 1))
        layers.extend([
            nn.ConvTranspose2d(latent_dim, out_features, 4, 1, 0, bias=False),
            nn.BatchNorm2d(out_features),
            nn.ReLU(True)
        ])
        
        # Upsampling blocks
        for i in range(num_blocks - 1):
            in_features = out_features
            out_features = in_features // 2
            layers.extend([
                nn.ConvTranspose2d(in_features, out_features, 4, 2, 1, bias=False),
                nn.BatchNorm2d(out_features),
                nn.ReLU(True)
            ])
            
        # Final layer: to img_size and target channels
        layers.extend([
            nn.ConvTranspose2d(out_features, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
        ])
        
        self.main = nn.Sequential(*layers)

    def forward(self, z):
        # Reshape z from (batch, latent_dim) to (batch, latent_dim, 1, 1)
        z = z.view(z.shape[0], -1, 1, 1)
        return self.main(z)

class Discriminator(nn.Module):
    """Powerful and FLEXIBLE DCGAN-style discriminator (Critic)"""
    def __init__(self, img_channels=3, img_size=64, feature_maps=64):
        super().__init__()
        
        # Calculate the number of downsampling blocks
        num_blocks = int(np.log2(img_size) - 2)
        if num_blocks < 1:
            raise ValueError(f"Image size {img_size} is too small for this DCGAN architecture. Minimum is 8.")

        layers = []
        
        # Initial layer
        layers.extend([
            nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        ])
        
        in_features = feature_maps
        # Downsampling blocks
        for i in range(num_blocks - 1):
            out_features = in_features * 2
            layers.extend([
                nn.Conv2d(in_features, out_features, 4, 2, 1, bias=False),
                nn.InstanceNorm2d(out_features, affine=True),
                nn.LeakyReLU(0.2, inplace=True)
            ])
            in_features = out_features
            
        # Final layer: downsample to a 1x1 feature map (a single score)
        layers.extend([
            nn.Conv2d(in_features, 1, 4, 1, 0, bias=False)
        ])
        
        self.main = nn.Sequential(*layers)

    def forward(self, img):
        return self.main(img)

class Discriminator(nn.Module):
    """Powerful and FLEXIBLE DCGAN-style discriminator (Critic)"""
    def __init__(self, img_channels=3, img_size=64, feature_maps=64):
        super().__init__()
        
        # Calculate the number of downsampling blocks
        num_blocks = int(np.log2(img_size) - 2)
        if num_blocks < 1:
            raise ValueError(f"Image size {img_size} is too small for this DCGAN architecture. Minimum is 8.")

        layers = []
        
        # Initial layer
        layers.extend([
            nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        ])
        
        in_features = feature_maps
        # Downsampling blocks
        for i in range(num_blocks - 1):
            out_features = in_features * 2
            layers.extend([
                nn.Conv2d(in_features, out_features, 4, 2, 1, bias=False),
                nn.InstanceNorm2d(out_features, affine=True),
                nn.LeakyReLU(0.2, inplace=True)
            ])
            in_features = out_features
            
        # Final layer: downsample to a 1x1 feature map (a single score)
        layers.extend([
            nn.Conv2d(in_features, 1, 4, 1, 0, bias=False)
        ])
        
        self.main = nn.Sequential(*layers)

    def forward(self, img):
        return self.main(img)

class Discriminator(nn.Module):
    """Powerful and FLEXIBLE DCGAN-style discriminator (Critic)"""
    def __init__(self, img_channels=3, img_size=64, feature_maps=64):
        super().__init__()
        
        # Calculate the number of downsampling blocks
        num_blocks = int(np.log2(img_size) - 2)
        if num_blocks < 1:
            raise ValueError(f"Image size {img_size} is too small for this DCGAN architecture. Minimum is 8.")

        layers = []
        
        # Initial layer
        layers.extend([
            nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True)
        ])
        
        in_features = feature_maps
        # Downsampling blocks
        for i in range(num_blocks - 1):
            out_features = in_features * 2
            layers.extend([
                nn.Conv2d(in_features, out_features, 4, 2, 1, bias=False),
                nn.InstanceNorm2d(out_features, affine=True),
                nn.LeakyReLU(0.2, inplace=True)
            ])
            in_features = out_features
            
        # Final layer: downsample to a 1x1 feature map (a single score)
        layers.extend([
            nn.Conv2d(in_features, 1, 4, 1, 0, bias=False)
        ])
        
        self.main = nn.Sequential(*layers)

    def forward(self, img):
        return self.main(img)

class Discriminator(nn.Module):
    """Powerful DCGAN-style discriminator (Critic)"""
    def __init__(self, img_channels=3, feature_maps=64):
        super().__init__()
        self.main = nn.Sequential(
            # Input size: (img_channels) x 64 x 64
            nn.Conv2d(img_channels, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps) x 32 x 32
            nn.Conv2d(feature_maps, feature_maps * 2, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 2, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*2) x 16 x 16
            nn.Conv2d(feature_maps * 2, feature_maps * 4, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 4, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*4) x 8 x 8
            nn.Conv2d(feature_maps * 4, feature_maps * 8, 4, 2, 1, bias=False),
            nn.InstanceNorm2d(feature_maps * 8, affine=True),
            nn.LeakyReLU(0.2, inplace=True),
            # State size: (feature_maps*8) x 4 x 4
            nn.Conv2d(feature_maps * 8, 1, 4, 1, 0, bias=False),
            # Final output is a single raw score, not a probability
        )

    def forward(self, img):
        return self.main(img)

class ImageDataset(Dataset):
    """Custom dataset for uploaded images"""
    def __init__(self, images, transform=None, img_size=64):
        self.images = images
        self.transform = transform
        self.img_size = img_size
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        image = self.images[idx]
        
        # Resize image
        image = image.resize((self.img_size, self.img_size), Image.Resampling.LANCZOS)
        
        if self.transform:
            image = self.transform(image)
        
        return image

def preprocess_images(uploaded_files, img_size=64):
    """Preprocess uploaded images"""
    images = []
    valid_files = 0
    
    progress_bar = st.progress(0, text="📄 Processing images...")
    status_text = st.empty()
    
    for idx, file in enumerate(uploaded_files):
        try:
            # Update progress with colorful text
            progress_bar.progress((idx + 1) / len(uploaded_files), 
                                text=f"🎨 Processing image {idx + 1}/{len(uploaded_files)}")
            status_text.text(f"📸 Current: {file.name}")
            
            image = Image.open(file)
            # Convert to RGB if needed
            if image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Basic validation - check if image is too small
            if image.size[0] < 16 or image.size[1] < 16:
                st.warning(f"⚠️ Skipping {file.name}: image too small (minimum 16x16)")
                continue
                
            images.append(image)
            valid_files += 1
            
        except Exception as e:
            st.error(f"❌ Could not process file {file.name}: {e}")
    
    progress_bar.empty()
    status_text.empty()
    
    return images, valid_files

def compute_gradient_penalty(discriminator, real_samples, fake_samples):
    """Calculates the gradient penalty for WGAN-GP"""
    alpha = torch.rand(real_samples.size(0), 1, 1, 1, device=device)
    interpolates = (alpha * real_samples + ((1 - alpha) * fake_samples)).requires_grad_(True)
    d_interpolates = discriminator(interpolates)
    fake = torch.ones(d_interpolates.size(), device=device, requires_grad=False)
    
    gradients = torch.autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=fake,
        create_graph=True,
        retain_graph=True,
        only_inputs=True,
    )[0]
    
    gradients = gradients.view(gradients.size(0), -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty

def train_simple_gan(images, epochs=50, batch_size=4, lr=0.0002, latent_dim=100, img_size=64):
    """Train the Simple GAN with basic loss"""
    
    # Data preparation
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    
    dataset = ImageDataset(images, transform=transform, img_size=img_size)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
    
    if len(dataloader) == 0:
        st.error(f"🚫 Cannot start training. The number of valid images ({len(dataset)}) is less than the batch size ({batch_size}). Please upload more images or reduce the batch size.")
        return None, None, [], []

    # Initialize networks
    generator = SimpleGenerator(latent_dim=latent_dim, img_size=img_size).to(device)
    discriminator = SimpleDiscriminator(img_size=img_size).to(device)
    
    # Loss and optimizers
    criterion = nn.BCELoss()
    optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))
    
    # Training progress containers with colorful headers
    st.markdown("### 🎯 Training Progress Dashboard")
    progress_bar = st.progress(0, text="🖼️ Starting training...")
    
    col1, col2, col3 = st.columns(3)
    with col1:
        d_loss_metric = st.metric("🔴 Discriminator Loss", "0.0000")
    with col2:
        g_loss_metric = st.metric("🔵 Generator Loss", "0.0000") 
    with col3:
        epoch_metric = st.metric("⏳ Current Epoch", "0")
    
    loss_chart = st.empty()
    sample_container = st.empty()
    
    # Training history
    d_losses = []
    g_losses = []
    
    # Training loop
    for epoch in range(epochs):
        epoch_d_loss = 0
        epoch_g_loss = 0
        batches = 0
        
        for i, real_images in enumerate(dataloader):
            batch_size_actual = real_images.size(0)
            real_images = real_images.to(device)
            
            # Labels
            real_labels = torch.ones(batch_size_actual, 1, device=device)
            fake_labels = torch.zeros(batch_size_actual, 1, device=device)
            
            # Train Discriminator
            optimizer_D.zero_grad()
            
            # Real images
            real_output = discriminator(real_images)
            d_loss_real = criterion(real_output, real_labels)
            
            # Fake images
            z = torch.randn(batch_size_actual, latent_dim, device=device)
            fake_images = generator(z)
            fake_output = discriminator(fake_images.detach())
            d_loss_fake = criterion(fake_output, fake_labels)
            
            d_loss = (d_loss_real + d_loss_fake) / 2
            d_loss.backward()
            optimizer_D.step()
            
            # Train Generator
            optimizer_G.zero_grad()
            
            fake_output = discriminator(fake_images)
            g_loss = criterion(fake_output, real_labels)
            g_loss.backward()
            optimizer_G.step()
            
            epoch_d_loss += d_loss.item()
            epoch_g_loss += g_loss.item()
            batches += 1
            
            # Memory cleanup
            del real_images, fake_images, z, real_output, fake_output
            if i % 5 == 0:  # Cleanup every 5 batches
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                gc.collect()
        
        if batches == 0:
            st.warning(f"⚠️ Epoch {epoch+1}/{epochs} - No batches processed. Is number of images < batch size?")
            continue
        
        # Record losses
        avg_d_loss = epoch_d_loss / batches
        avg_g_loss = epoch_g_loss / batches
        d_losses.append(avg_d_loss)
        g_losses.append(avg_g_loss)
        
        # Update progress and metrics with emojis
        progress = (epoch + 1) / epochs
        progress_bar.progress(progress, text=f"🎨 Training in progress... Epoch {epoch+1}/{epochs}")
        
        d_loss_metric.metric("🔴 Discriminator Loss", f"{avg_d_loss:.4f}")
        g_loss_metric.metric("🔵 Generator Loss", f"{avg_g_loss:.4f}")
        epoch_metric.metric("⏳ Current Epoch", f"{epoch+1}/{epochs}")
        
        # Update loss chart every 5 epochs with colorful styling
        if (epoch + 1) % 5 == 0 or epoch == epochs - 1:
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.plot(d_losses, label='🔴 Discriminator Loss', color='#ff6b6b', linewidth=2.5, alpha=0.8)
            ax.plot(g_losses, label='🔵 Generator Loss', color='#4ecdc4', linewidth=2.5, alpha=0.8)
            ax.set_xlabel('Epoch', fontsize=12, fontweight='bold')
            ax.set_ylabel('Loss', fontsize=12, fontweight='bold')
            ax.set_title('📈 Training Loss Curves', fontsize=14, fontweight='bold', pad=20)
            ax.legend(fontsize=11)
            ax.grid(True, alpha=0.3, linestyle='--')
            ax.set_facecolor('#f8f9fa')
            fig.patch.set_facecolor('white')
            loss_chart.pyplot(fig)
            plt.close(fig)
        
        # Generate sample images every 10 epochs with better styling
        if (epoch + 1) % 10 == 0 or epoch == epochs - 1:
            with torch.no_grad():
                sample_z = torch.randn(4, latent_dim, device=device)
                sample_images = generator(sample_z)
                sample_images = (sample_images + 1) / 2  # Denormalize
                
                fig, axes = plt.subplots(1, 4, figsize=(14, 4))
                fig.suptitle(f'🎨 Generated Images - Epoch {epoch+1}', fontsize=16, fontweight='bold', y=1.05)
                
                for j in range(4):
                    img = sample_images[j].cpu().permute(1, 2, 0).numpy()
                    img = np.clip(img, 0, 1)
                    axes[j].imshow(img)
                    axes[j].axis('off')
                    axes[j].set_title(f'✨ Sample {j+1}', fontsize=12, pad=10)
                    # Add colorful border
                    for spine in axes[j].spines.values():
                        spine.set_edgecolor(['#ff6b6b', '#4ecdc4', '#45b7d1', '#96ceb4'][j])
                        spine.set_linewidth(3)
                
                plt.tight_layout()
                sample_container.pyplot(fig)
                plt.close(fig)
                
                del sample_z, sample_images
    
    return generator, discriminator, d_losses, g_losses

def train_dcgan_wgan(images, epochs=50, batch_size=4, lr=0.0002, latent_dim=100, img_size=64):
    """Train the WGAN-GP with memory-efficient approach and data augmentation"""


    if img_size != 64:
        
        st.warning(f"⚠️ DCGAN-WGAN model is optimized for 64x64 images. Forcing resolution to 64px.")
        img_size = 64
    
    #st.warning(f"⚠️ DCGAN-WGAN model is optimized for 64x64 images. Forcing resolution to 64px.")
    
    # 1. HEAVY DATA AUGMENTATION
    # This is key for training on a small number of images
    transform = transforms.Compose([
        transforms.RandomResizedCrop(img_size, scale=(0.8, 1.0)),
        transforms.RandomHorizontalFlip(),
        transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        transforms.RandomRotation(15),
        transforms.ToTensor(),
        transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])
    ])
    
    dataset = ImageDataset(images, transform=transform, img_size=img_size)
    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True)
    
    if len(dataloader) == 0:
        st.error(f"🚫 Cannot start training. The number of valid images ({len(dataset)}) is less than the batch size ({batch_size}). Please upload more images or reduce the batch size.")
        return None, None, [], []

    # 2. UPDATED MODEL INITIALIZATION (PASSING IMG_SIZE)
    # This now uses the flexible models which adapt to the image size
    st.info(f"🖼️ Initializing DCGAN-WGAN model for {img_size}x{img_size} images...")
    if img_size == 64:
        generator = Generator(latent_dim=latent_dim, img_channels=3, feature_maps=64).to(device)
        discriminator = Discriminator(img_channels=3, feature_maps=64).to(device)
    elif img_size == 128:
        generator = Generator128(latent_dim=latent_dim, img_channels=3, feature_maps=64).to(device)
        discriminator = Discriminator128(latent_dim=latent_dim, img_channels=3, feature_maps=64).to(device)
    else:
    # This will handle 32px or any other unsupported sizes for this model
        st.error(f"🚫 The DCGAN-WGAN model only supports 64x64 or 128x128 resolution.")
    # WGAN-GP parameters
    lambda_gp = 10
    n_critic = 5 # Train discriminator more often than generator
    
    # Use Adam optimizers with parameters recommended for WGAN
    optimizer_G = optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.9))
    optimizer_D = optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.9))
    
    st.markdown("### 🎯 Training Progress Dashboard")
    progress_bar = st.progress(0, text="🖼️ Starting training...")
    
    col1, col2, col3 = st.columns(3)
    with col1:
        d_loss_metric = st.metric("🔴 Critic Loss", "0.00")
    with col2:
        g_loss_metric = st.metric("🔵 Generator Loss", "0.00") 
    with col3:
        epoch_metric = st.metric("⏳ Current Epoch", "0")
    
    loss_chart = st.empty()
    sample_container = st.empty()
    
    d_losses = []
    g_losses = []
    
    # 3. NEW WGAN-GP TRAINING LOOP
    for epoch in range(epochs):
        epoch_d_loss = 0
        epoch_g_loss = 0
        batches = 0
        
        for i, real_images in enumerate(dataloader):
            real_images = real_images.to(device)
            batch_size_actual = real_images.size(0)
            
            # ---------------------
            #  Train Discriminator (Critic)
            # ---------------------
            optimizer_D.zero_grad()
            
            # Sample noise as generator input
            z = torch.randn(batch_size_actual, latent_dim, device=device)
            fake_images = generator(z)
            
            # Real and fake images scores
            real_validity = discriminator(real_images)
            fake_validity = discriminator(fake_images.detach())
            
            # Gradient penalty
            gradient_penalty = compute_gradient_penalty(discriminator, real_images.data, fake_images.data)
            
            # Critic loss
            d_loss = -torch.mean(real_validity) + torch.mean(fake_validity) + lambda_gp * gradient_penalty
            
            d_loss.backward()
            optimizer_D.step()
            
            epoch_d_loss += d_loss.item()
            
            # Train the generator only every n_critic iterations
            if i % n_critic == 0:
                # -----------------
                #  Train Generator
                # -----------------
                optimizer_G.zero_grad()
                
                # Generate a batch of images
                # Reuse the z from the critic step for efficiency
                gen_imgs = generator(z)
                
                # Generator loss
                g_loss = -torch.mean(discriminator(gen_imgs))
                
                g_loss.backward()
                optimizer_G.step()
                
                epoch_g_loss += g_loss.item()

            batches += 1
            
            # Memory cleanup
            del real_images, fake_images, z
            if torch.cuda.is_available():
                torch.cuda.empty_cache()

        if batches == 0:
            st.warning(f"⚠️ Epoch {epoch+1}/{epochs} - No batches processed.")
            continue
        
        avg_d_loss = epoch_d_loss / batches
        avg_g_loss = epoch_g_loss / (batches / n_critic) if batches > 0 else 0
        d_losses.append(avg_d_loss)
        g_losses.append(avg_g_loss)
        
        progress = (epoch + 1) / epochs
        progress_bar.progress(progress, text=f"🎨 Training in progress... Epoch {epoch+1}/{epochs}")
        
        d_loss_metric.metric("🔴 Critic Loss", f"{avg_d_loss:.2f}")
        g_loss_metric.metric("🔵 Generator Loss", f"{avg_g_loss:.2f}")
        epoch_metric.metric("⏳ Current Epoch", f"{epoch+1}/{epochs}")
        
        # Update loss chart every 5 epochs with colorful styling
        if (epoch + 1) % 5 == 0 or epoch == epochs - 1:
            fig, ax = plt.subplots(figsize=(10, 5))
            ax.plot(d_losses, label='🔴 Critic Loss', color='#ff6b6b', linewidth=2.5, alpha=0.8)
            ax.plot(g_losses, label='🔵 Generator Loss', color='#4ecdc4', linewidth=2.5, alpha=0.8)
            ax.set_xlabel('Epoch', fontsize=12, fontweight='bold')
            ax.set_ylabel('Loss', fontsize=12, fontweight='bold')
            ax.set_title('📈 Training Loss Curves', fontsize=14, fontweight='bold', pad=20)
            ax.legend(fontsize=11)
            ax.grid(True, alpha=0.3, linestyle='--')
            ax.set_facecolor('#f8f9fa')
            fig.patch.set_facecolor('white')
            loss_chart.pyplot(fig)
            plt.close(fig)

        # Generate sample images every 10 epochs with better styling
        if (epoch + 1) % 10 == 0 or epoch == epochs - 1:
            with torch.no_grad():
                sample_z = torch.randn(4, latent_dim, device=device)
                sample_images = generator(sample_z)
                sample_images = (sample_images + 1) / 2
                
                fig, axes = plt.subplots(1, 4, figsize=(14, 4))
                fig.suptitle(f'🎨 Generated Images - Epoch {epoch+1}', fontsize=16, fontweight='bold', y=1.05)
                
                for j in range(4):
                    img = sample_images[j].cpu().permute(1, 2, 0).numpy()
                    img = np.clip(img, 0, 1)
                    axes[j].imshow(img)
                    axes[j].axis('off')
                    axes[j].set_title(f'✨ Sample {j+1}', fontsize=12, pad=10)
                    for spine in axes[j].spines.values():
                        spine.set_edgecolor(['#ff6b6b', '#4ecdc4', '#45b7d1', '#96ceb4'][j])
                        spine.set_linewidth(3)
                
                plt.tight_layout()
                sample_container.pyplot(fig)
                plt.close(fig)
                
                del sample_z, sample_images
                
    return generator, discriminator, d_losses, g_losses

def generate_images(generator, num_images=8, latent_dim=100, is_simple=False):
    """Generate new images using trained generator"""
    generator.eval()
    with torch.no_grad():
        z = torch.randn(num_images, latent_dim, device=device)
        if not is_simple:
            # For DCGAN, reshape z for ConvTranspose2d input
            z = z.view(z.shape[0], -1, 1, 1)
        generated_images = generator(z)
        generated_images = (generated_images + 1) / 2  # Denormalize
        
        # Create grid with better styling
        rows = 2
        cols = (num_images + 1) // rows
        fig, axes = plt.subplots(rows, cols, figsize=(15, 8))
        axes = axes.flatten()
        
        colors = ['#ff6b6b', '#4ecdc4', '#45b7d1', '#96ceb4', '#feca57', '#ff9ff3', '#54a0ff', '#5f27cd']
        
        for i in range(num_images):
            img = generated_images[i].cpu().permute(1, 2, 0).numpy()
            img = np.clip(img, 0, 1)
            axes[i].imshow(img)
            axes[i].axis('off')
            axes[i].set_title(f'✨ Generated #{i+1}', fontsize=11, fontweight='bold', pad=8)
            # Add colorful border
            for spine in axes[i].spines.values():
                spine.set_edgecolor(colors[i % len(colors)])
                spine.set_linewidth(3)
        
        # Turn off unused subplots
        for i in range(num_images, len(axes)):
            axes[i].axis('off')
        
        fig.suptitle('🎨 AI Generated Masterpieces', fontsize=18, fontweight='bold', y=0.98)
        fig.patch.set_facecolor('white')
        plt.tight_layout()
        return fig


# ===================== DATABASE SETUP =====================
# Initialize TinyDB. This will create a db.json file in the same directory.
db = TinyDB('models_db.json')

def save_model_to_db(name, generator, model_type, latent_dim, img_size):
    """Saves a model's state_dict and metadata to TinyDB."""
    st.info(f"💾 Saving model '{name}' to database...")
    
    # Save model state_dict to an in-memory buffer
    buffer = io.BytesIO()
    torch.save(generator.state_dict(), buffer)
    buffer.seek(0)
    
    # Encode the binary data to a base64 string
    model_b64 = base64.b64encode(buffer.read()).decode('utf-8')
    
    # Create the document to insert into TinyDB
    model_doc = {
        "name": name,
        "model_type": model_type,
        "latent_dim": latent_dim,
        "img_size": img_size,
        "timestamp": str(datetime.datetime.now()),
        "state_dict_b64": model_b64
    }
    
    # Insert or update the model by name
    ModelQuery = Query()
    db.upsert(model_doc, ModelQuery.name == name)
    st.success(f"✅ Model '{name}' saved successfully!")


def load_model_from_db(name):
    """Loads a model's state_dict from TinyDB and reconstructs the model."""
    st.info(f"🔍 Loading model '{name}' from database...")
    
    # Find the model document
    ModelQuery = Query()
    model_doc = db.get(ModelQuery.name == name)
    
    if not model_doc:
        st.error(f"Model '{name}' not found in the database.")
        return None, None

    # Get metadata
    model_type = model_doc['model_type']
    latent_dim = model_doc['latent_dim']
    img_size = model_doc['img_size']
    
    # Instantiate the correct model architecture
    if model_type == 'Simple GAN':
        generator = SimpleGenerator(latent_dim=latent_dim, img_size=img_size).to(device)
    elif model_type == 'DCGAN-WGAN':
        if img_size == 128:
            generator = Generator128(latent_dim=latent_dim).to(device)
        else: # Default to 64px DCGAN
            generator = Generator(latent_dim=latent_dim).to(device)
    else:
        st.error(f"Unknown model type: {model_type}")
        return None, None
        
    # Decode the base64 string back to binary data
    model_b64 = model_doc['state_dict_b64']
    model_bytes = base64.b64decode(model_b64.encode('utf-8'))
    
    # Load the state_dict from the in-memory buffer
    buffer = io.BytesIO(model_bytes)
    generator.load_state_dict(torch.load(buffer, map_location=device))
    
    st.success(f"✅ Model '{name}' loaded and ready!")
    return generator, model_doc



# Streamlit UI
def main():
    # Colorful animated title
    st.markdown('<h1 class="big-title"> AI Art Generator Studio</h1>', unsafe_allow_html=True)
    
    # Colorful subtitle with gradient
    st.markdown("""
    <div class="main-subtitle" style="text-align: center; margin-bottom: 2rem;">
        <h3 style="font-weight: 300;">
            ✨ Train powerful GANs on your images • 🖼️ Multiple model architectures 
        </h3>
    </div>
    """, unsafe_allow_html=True)
    
    # Sidebar with colorful styling
    with st.sidebar:
        st.markdown("###  **Training Control Panel**")
        
        # Add some colorful info boxes
        st.info("💡 **Pro Tip**: DCGAN-WGAN works better with fewer images!")
        
        # MODEL SELECTION DROPDOWN
        st.markdown("#### 🦄 **Model Architecture**")
        model_type = st.selectbox(
            "🧠 Choose GAN Model",
            ["Simple GAN", "DCGAN-WGAN"],
            index=0,
            help="Simple GAN: Fast, good for beginners | DCGAN-WGAN: Advanced, better quality"
        )
        
        # Show model info
        if model_type == "Simple GAN":
            st.success("⚡ **Fast Training** • Good for beginners")
        else:
            st.success("🎨 **High Quality** • Advanced architecture")
        
        st.markdown("#### 🖼️ **Image Settings**")
        img_size = st.selectbox("📏 Image Resolution", [32, 64, 128], index=1, 
                               help="Higher = better quality, more memory")
        
        st.markdown("#### ⚡ **Training Parameters**")
        epochs = st.slider("🔥 Training Epochs", 10, 2000, 500, 
                          help="More epochs = better results, longer training")
        batch_size = st.slider("📦 Batch Size", 1, 8, 4,
                              help="Lower if running out of memory")
        lr = st.number_input("🎯 Learning Rate", 0.0001, 0.01, 0.0002, format="%.4f",
                           help="Controls how fast the AI learns")
        latent_dim = st.slider("🌌 Latent Dimension", 50, 200, 100,
                              help="Complexity of generated features")
        
        # Memory warning with emoji
        if img_size > 64 or batch_size > 4:
            st.warning("⚠️ **Memory Alert**: High settings may exceed 4GB limit!")
        else:
            st.success("✅ **Memory**: Settings look good!")
        
        # Device info with colors
        device_emoji = "🖼️" if device.type == "cuda" else "💻"
        st.info(f"{device_emoji} **Device**: {device.type.upper()}")
        
        st.markdown("---")
        st.markdown("### 📤 **Upload Your Training Images**")
        
        # Styled file uploader
        uploaded_files = st.file_uploader(
            "🖼️ Choose your images (1-100)", 
            type=['png', 'jpg', 'jpeg'], 
            accept_multiple_files=True,
            help="Upload 1-100 images to train your AI artist! 🎨"
        )
        
        if uploaded_files:
            num_files = len(uploaded_files)
            if num_files > 100:
                st.sidebar.error("🚫 **Too many files!** Please upload maximum 100 images")
                return
            
            # Progress bar for file count
            progress_percentage = min(num_files / 20, 1.0)  # Ideal around 20 images
            st.sidebar.progress(progress_percentage, f"📈 Dataset completeness: {num_files}/20+ images")
            
            # Show sample images with better styling
            if st.sidebar.checkbox(" **Preview uploaded images**", help="See what you're training on"):
                st.sidebar.markdown("#### 🖼️ **Image Preview**")
                num_previews = min(4, num_files)
                for i, file in enumerate(uploaded_files[:num_previews]):
                    try:
                        img = Image.open(file)
                        st.sidebar.image(img, caption=f"✨ Image {i+1}", use_container_width=True)
                    except Exception as e:
                        st.sidebar.error(f"❌ Error: {file.name}")
    
    # Main interface with tabs
    tab1, tab2, tab3, tab4 = st.tabs(["📤 **Upload & Train**", "🎨 **Generate Art**", "📊 **Model Info**", "🖼️ **Eye Friendly**"])
    
    with tab1:
        # Centered training control since upload is in the sidebar
        _, center_col, _ = st.columns([1, 2, 1])
        
        with center_col:
            st.markdown("### 🖼️ **Start Training**")
            
            if not uploaded_files:
                st.info(" **Upload images in the sidebar first** to start training your AI!")
            else:
                # Memory estimation with colors
                estimated_memory = (len(uploaded_files) * img_size * img_size * 3 * 4) / (1024**3)
                
                col_mem1, col_mem2 = st.columns(2)
                with col_mem1:
                    st.metric("💾 **Est. Memory**", f"{estimated_memory:.1f} GB")
                with col_mem2:
                    memory_status = "🟢 Good" if estimated_memory < 2 else "🟡 High" if estimated_memory < 3 else "🔴 Too High"
                    st.metric("📊 **Memory Status**", memory_status)
                
                if estimated_memory > 3:
                    st.warning("⚠️ **High memory usage!** Consider reducing image size or batch size.")
                
                # Display selected model info
                st.info(f"🦄 **Selected Model**: {model_type}")
                
                # Big colorful training button
                if st.button("🖼️ **START TRAINING**", type="primary", use_container_width=True):
                    if len(uploaded_files) < 1:
                        st.error("🚫 Please upload at least 1 image")
                        return
                    
                    with st.spinner("🔥 Preprocessing your masterpieces..."):
                        images, valid_files = preprocess_images(uploaded_files, img_size)
                    
                    if valid_files == 0:
                        st.error("❌ No valid images found")
                        return
                    
                    # Celebration animation
                    st.success(f"✅ Processing {valid_files} beautiful images")
                    
                    # Train the GAN based on selected model
                    try:
                        if model_type == "Simple GAN":
                            generator, discriminator, d_losses, g_losses = train_simple_gan(
                                images, epochs, batch_size, lr, latent_dim, img_size
                            )
                        else:  # DCGAN-WGAN
                            generator, discriminator, d_losses, g_losses = train_dcgan_wgan(
                                images, epochs, batch_size, lr, latent_dim, img_size
                            )
                        
                        if generator is not None:
                            # Another celebration
                            st.success("🎉 **Training Complete!** Your AI artist is ready!")
                            
                            # Store models in session state
                            st.session_state['generator'] = generator
                            st.session_state['discriminator'] = discriminator
                            st.session_state['trained'] = True
                            st.session_state['latent_dim'] = latent_dim
                            st.session_state['model_type'] = model_type
                            
                            # Show training summary
                            col_summary1, col_summary2 = st.columns(2)
                            with col_summary1:
                                loss_label = "🎯 **Final D-Loss**" if model_type == "Simple GAN" else "🎯 **Final C-Loss**"
                                st.metric(loss_label, f"{d_losses[-1]:.4f}")
                            with col_summary2:
                                st.metric("🎨 **Final G-Loss**", f"{g_losses[-1]:.4f}")
                    
                    except RuntimeError as e:
                        if "out of memory" in str(e).lower():
                            st.error("🔥 **Out of memory!** Try reducing batch size or image size.")
                        else:
                            st.error(f"💥 Training error: {e}")
                    except Exception as e:
                        st.error(f"⚡ Unexpected error: {e}")
        # This UI will now appear AFTER training is complete and will stay visible
        if st.session_state.get('trained', False):
            st.markdown("---")
            st.markdown("### 💾 **Save Your Trained Model**")
            model_name = st.text_input("Enter a name for your model (e.g., 'Abstract Art v1')")
            if st.button("💾 **SAVE MODEL TO DATABASE**", use_container_width=True):
                if model_name:
                    # We grab the generator and parameters FROM session_state
                    save_model_to_db(
                        name=model_name,
                        generator=st.session_state['generator'],
                        model_type=st.session_state['model_type'],
                        latent_dim=st.session_state['latent_dim'],
                        img_size=img_size # img_size is still available from the sidebar
                    )
                else:
                    st.warning("⚠️ Please enter a name for the model before saving.")
    
    
    with tab2:
        st.markdown("### 🎨 **Generate Amazing Art**")
        
        # Get all saved models from the database
        saved_models = db.all()
        model_names = [model['name'] for model in saved_models]

        # This `if` statement now checks if there are NO saved models AND a model has NOT been trained in the current session.
        # This replaces the old `else` block's functionality.
        if not model_names and not st.session_state.get('trained', False):
            st.info("👈 **Train and save your first AI model** in the 'Upload & Train' tab to start generating art!")
            
            # Show some example placeholder
            st.markdown("#### 🌟 **What you'll be able to create:**")
            st.markdown("""
            - 🖼️ **Unique AI-generated images** based on your training data
            - 🎨 **Infinite variations** with different random seeds  
            - 📊 **Batch generation** of multiple images at once
            - 🎭 **Style consistency** learned from your uploaded images
            - 🦄 **Choice of model architectures** for different quality levels
            """)
        else:
            # This part runs if you HAVE saved models or just finished training one.
            col_gen1, col_gen2 = st.columns([1, 2], gap="large")
            
            with col_gen1:
                st.markdown("#### 🧠 **Select Your AI Artist**")

                # Create a list of model choices for the dropdown
                display_names = list(model_names)
                # Add the currently trained model as the first option if it exists
                if st.session_state.get('trained', False):
                    display_names.insert(0, "✨ Most Recently Trained Model")

                selected_model_name = st.selectbox(
                    "Choose a model to use for generation:",
                    options=display_names
                )

                if selected_model_name:
                    generator_to_use = None
                    model_info = {}

                    # Logic to decide which generator to use
                    if selected_model_name == "✨ Most Recently Trained Model":
                        st.info("Using the model you just trained in this session.")
                        generator_to_use = st.session_state.get('generator')
                        model_info['latent_dim'] = st.session_state.get('latent_dim')
                        model_info['model_type'] = st.session_state.get('model_type')
                    else:
                        # Load the selected model from the database
                        generator_to_use, model_info = load_model_from_db(selected_model_name)

                    if generator_to_use and model_info:
                        st.markdown("#### 🎛️ **Generation Settings**")
                        num_generate = st.slider("🖼️ **Number of images**", 1, 16, 8)
                        
                        st.markdown("#### 🎲 **Random Seed**")
                        use_seed = st.checkbox("🔒 **Use fixed seed**", help="For reproducible results")
                        if use_seed:
                            seed_value = st.number_input("🌱 **Seed value**", 0, 9999, 42)
                            torch.manual_seed(seed_value)
                        
                        if st.button("🎨 **GENERATE ART**", type="primary", use_container_width=True):
                            with st.spinner("🎭 Your AI is painting masterpieces..."):
                                try:
                                    is_simple = model_info.get('model_type', 'Simple GAN') == 'Simple GAN'
                                    fig = generate_images(
                                        generator_to_use, 
                                        num_generate, 
                                        model_info.get('latent_dim', 100),
                                        is_simple=is_simple
                                    )
                                    with col_gen2:
                                        st.pyplot(fig)
                                        st.success("🎉 **Fresh art created!**")
                                    plt.close(fig)
                                except Exception as e:
                                    st.error(f"🎨 Generation error: {e}")
            
            with col_gen2:
                st.info("🎭 **Generated images will appear here** after you click the generate button!")
    
    with tab3:
        st.markdown("### 📊 **Model Architecture & Info**")
        
        if st.session_state.get('trained', False):
            current_model = st.session_state.get('model_type', 'Unknown')
            
            col_info1, col_info2 = st.columns(2)
            
            with col_info1:
                st.markdown("#### 🧠 **Generator Network**")
                st.success("✅ **Status**: Trained and Ready")
                st.info(f"🦄 **Model Type**: {current_model}")
                st.info(f"🌌 **Latent Dimensions**: {st.session_state.get('latent_dim', 100)}")
                
                if current_model == "Simple GAN":
                    st.info("🏗️ **Architecture**: Lightweight CNN with Upsampling")
                else:
                    st.info("🏗️ **Architecture**: DCGAN with ConvTranspose2d")
                
            with col_info2:
                st.markdown("#### 🕵️ **Discriminator Network**") 
                st.success("✅ **Status**: Trained and Ready")
                st.info("🔍 **Purpose**: Real vs Fake Image Classification")
                
                if current_model == "Simple GAN":
                    st.info("🏗️ **Architecture**: Convolutional Classifier")
                else:
                    st.info("🏗️ **Architecture**: WGAN-GP Critic")
                    st.info("⚖️ **Loss**: Wasserstein + Gradient Penalty")
                
        else:
            st.info("📊 **Model information will appear here after training**")
            
        st.markdown("#### 🔧 **Technical Specifications**")
        tech_col1, tech_col2, tech_col3 = st.columns(3)
        
        with tech_col1:
            st.metric("🖥️ **Framework**", "PyTorch")
        with tech_col2:  
            st.metric("💾 **Memory Optimized**", "4GB")
        with tech_col3:
            current_arch = st.session_state.get('model_type', 'Not Selected')
            st.metric("⚡ **Architecture**", current_arch)
        
        # Model comparison table
        st.markdown("#### 📈 **Model Comparison**")
        comparison_data = {
            "Feature": ["Training Speed", "Image Quality", "Memory Usage", "Stability", "Data Augmentation"],
            "Simple GAN": ["⚡ Fast", "🟡 Good", "💚 Low", "🟡 Moderate", "❌ None"],
            "DCGAN-WGAN": ["🐌 Slower", "💚 Excellent", "🟡 Higher", "💚 Very Stable", "✅ Heavy"]
        }
        st.table(comparison_data)

    with tab4:
        st.markdown("### 🖼️ **Eye Friendly UI**")
        st.info("Toggle the switch below to slightly blur the highly-detailed background image, making the text easier to read for extended periods.")
        
        # Checkbox to toggle blur and save space in session state
        blur_bg = st.checkbox("🌫️ **Slightly Blur Background**", value=st.session_state.get('blur_bg_state', False), key='blur_bg_checkbox')
        st.session_state['blur_bg_state'] = blur_bg
        
        if blur_bg:
            st.markdown("""
            <style>
                /* Blur the main container, header, and sidebar to affect the background underneath */
                [data-testid="stAppViewContainer"], 
                [data-testid="stSidebar"], 
                [data-testid="stHeader"] {
                    backdrop-filter: blur(10px) !important;
                    -webkit-backdrop-filter: blur(10px) !important;
                }
                
                /* Ensure children of sidebar don't double-blur */
                [data-testid="stSidebar"] > div:first-child {
                    backdrop-filter: none !important;
                }
            </style>
            """, unsafe_allow_html=True)
            st.success("Background blur activated! Text is now easier to read.")
    
    # Bottom section with cleanup
    st.markdown("---")
    col_clean1, col_clean2, col_clean3 = st.columns([2, 1, 1])
    
    with col_clean2:
        if st.button("🧹 **Clear Memory**", help="Clear GPU/CPU memory and reset app state", use_container_width=True):
            # Clear specific keys
            for key in ['generator', 'discriminator', 'trained', 'latent_dim', 'model_type']:
                if key in st.session_state:
                    del st.session_state[key]
            
            # Clear GPU cache
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            gc.collect()
            
            st.success("🧹 Memory cleared and state reset!")
    
    with col_clean3:
        if st.button("ℹ️ **About**", use_container_width=True):
            st.info("""
            🎨 **AI Art Generator Studio**
            
            A comprehensive GAN training application featuring:
            • 💾 4GB memory optimization
            • 🖼️ Multiple model architectures
            • 🎯 High-quality results
            • 🖼️ Custom image datasets
            • 📊 Advanced training techniques
            
            **Models Available:**
            • Simple GAN: Fast training, good for beginners
            • DCGAN-WGAN: Advanced architecture with better quality
            """)

    # --- NEW THEORY & MATH SECTION ---
    st.markdown("---")
    with st.expander("📚 **AI Theory & Deep Learning Deep-Dive**", expanded=True):
        t_col1, t_col2 = st.columns([1, 1], gap="medium")
        
        with t_col1:
            st.markdown("### 🧬 **I. The Adversarial Game & Nash Equilibrium**")
            st.markdown("""
            Generative Adversarial Networks (GANs) represent a paradigm shift from traditional supervised learning to unsupervised density estimation via a competitive game. The training process is not a standard optimization toward a global minimum, but a search for a **Nash Equilibrium** in a zero-sum game between two neural networks. In this landscape, the loss surface is non-convex and highly dynamic, as the "target" for the Generator is constantly moving.
            
            **The Training Tension:**
            In the classic formulation by Ian Goodfellow (2014), the Generator (G) and Discriminator (D) engage in a minimax game. If D becomes too powerful too quickly, G's gradients "saturate" or vanish, and learning plateaus. Conversely, if G is too strong, it may exploit specific mathematical weaknesses in D rather than learning the generalized data distribution. This implementation uses Non-Saturating Loss heuristics and Wasserstein metrics to ensure that G always has a "signal" to follow, even when D is highly accurate.
            
            **Game Theory Dynamics:**
            The convergence in GANs is often oscillatory. Because one network's gain is the other's loss, they can end up chasing each other in circles in the parameter space (known as limit cycles). We stabilize this using Adam's second-order momentum and low beta-1 values to "dampen" these oscillations, ensuring the networks settle into a stable state where the generated art mirrors the training data's underlying probability distribution.
            """)
            
            st.markdown("### 🌌 **II. Latent Space Manifolds & Disentanglement**")
            st.markdown("""
            The latent noise vector (z) is the "seed" for a **non-linear manifold traversal**. The Generator learns to map a simple distribution (usually a 100-dimensional Gaussian) to the complex, high-dimensional space of natural image pixels. This mapping effectively "unfolds" the high-dimensional data manifold into a simpler, more navigable space.
            
            **Manifold Continuity:**
            In a well-trained GAN, the latent space is continuous. This means if you take two random points, A and B, in the 100D latent space and slowly interpolate between them, the resulting images will morph smoothly. In deep learning theory, this is known as **Latent Factor Disentanglement**—where individual dimensions of z start to correspond to independent features like "texture," "perspective," or "color palette," rather than being entangled in a noisy mess.
            
            **Vector Arithmetic in Latent Space:**
            An amazing property of these manifolds is that you can perform "visual arithmetic." While not explicitly coded in this app, a perfectly disentangled latent space would allow you to take the vector for "Blue Sky," subtract the vector for "Daytime," and add the vector for "Night," resulting in a "Night Sky" generation. This demonstrates that the AI has internalised a structured understanding of visual concepts.
            """)

            st.markdown("### ⚡ **III. Activation Functions: LeakyReLU vs. ReLU**")
            st.markdown("""
            Standard ReLU (Rectified Linear Unit) functions are "killing" neurons. If a neuron's value drops below zero, its gradient becomes zero, and it never learns again—a phenomenon called the "Dying ReLU" problem.
            
            **The LeakyReLU Solution:**
            To maintain the "dreamy" complexity of GAN art, we use **LeakyReLU** in the Discriminator with a leak slope of 0.2. This ensures that even when a neuron is inactive, it still passes a small gradient (0.2 * input). This "leak" is vital for the early stages of training when the Discriminator is seeing mostly noise; it allows almost all neurons to contribute to the initial gradient flow, helping the AI find the "rough sketch" of your images much faster than standard architectures.
            """)

            st.markdown("### ⚠️ **IV. Mode Collapse & Diversity Metrics**")
            st.markdown("""
            The most significant technical hurdle in GANs is **Mode Collapse**. This occurs when the Generator discovers a small subset of the training data (a "mode") that consistently fools the Discriminator. Instead of learning to generate the entire variety of your art, the Generator "collapses" and starts producing the exact same image (or a narrow set of variations) repeatedly.
            
            **Mitigation Strategies:**
            This app utilizes several techniques to combat this:
            - **Batch Normalization/Instance Normalization:** These layers help stabilize the internal covariate shift, preventing the networks from becoming "locked" into a single output mode.
            - **Heavy Data Augmentation:** By rotating and cropping images randomly, we force the Discriminator to be more generalized, which in turn forces the Generator to be more diverse to succeed.
            - **Mini-batch Discrimination:** While technically complex, the idea is to let the Discriminator look at a whole batch of images at once. If every image in G's batch looks identical, D can penalize G for lack of diversity, forcing it to branch out into new artistic styles.
            """)

            st.markdown("### 📏 **V. Earth Mover Theory & Wasserstein-1**")
            st.markdown("""
            Standard GANs use the Jensen-Shannon (JS) divergence as a loss function, which is mathematically "brittle" in high dimensions. If the Generator's distribution and the Real data distribution do not overlap significantly, the JS gradient becomes exactly zero, and training dies. This is why early GANs were so difficult to tune.
            
            **The Earth Mover Advantage:**
            WGAN-GP (Wasserstein GAN with Gradient Penalty) replaces JS divergence with the **Earth Mover's Distance**. Imagine the two probability distributions as piles of sand. The Wasserstein metric measures the minimum "work" (mass * distance) required to transport the sand from one pile to match the shape of the other. Because this distance is continuous and differentiable almost everywhere, it provides the smooth, non-vanishing "Critic Loss" curves we see in the dashboard, even when the Generator is still producing total noise.
            """)

        with t_col2:
            st.markdown("### 🎨 **VI. Non-Saturating Gradient Flows**")
            st.markdown("""
            A critical technical challenge in training "Adversarial" networks is ensuring the gradient signal never fully disappears. In the early stages, the Discriminator can easily distinguish between noise and art. If G minimizes the probability of D being correct, the gradient "saturates" (flattens out) near zero because D is too certain.
            
            **Practical Optimization:**
            Instead, we use a "Non-Saturating" heuristic where G maximizes the probability of D being wrong. This leads to much larger gradients in the early "birth" phase of your AI's artistic career. This ensures that the weights in the deep convolutional layers are updated with enough force to pull meaningful features out of the initial Gaussian noise distribution, preventing the "dead start" bug common in simpler GAN setups.
            """)

            st.markdown("### 🏗️ **VII. Transposed Convolution Dynamics**")
            st.markdown("""
            The Generator works the opposite of an image classifier. Instead of "pooling" an image down to a single label, it uses **Fractionally-strided Convolutions** (Transposed Convolutions) to "expand" the 1x1 latent vector into a high-resolution pixel grid. These are sometimes called "Deconvolutional" layers, though that term is technically a misnomer.
            
            **Checkerboard Artifact Prevention:**
            Visible grid patterns or "checkerboard artifacts" are a common failure in GANs, caused by uneven overlapping of kernels during up-sampling. We mitigate this by using specific stride and kernel size ratios (typically 4x4 kernels with stride 2). This ensures that every pixel is recalculated with the same frequency by the neural "brush," resulting in the smooth, organic textures found in high-quality GAN outputs rather than the mechanical "window-pane" effect.
            """)

            st.markdown("### ⚖️ **VIII. Lipschitz Continuity & Penalty**")
            st.markdown("""
            For the Wasserstein loss to remain valid, the Critic network must be an **"almost everywhere" differentiable 1-Lipschitz function**. This means the slope of the Critic's output can never exceed 1. Early WGANs tried to enforce this by "Weight Clipping," which severely limited the network's capacity and caused its weights to explode or vanish as they were forced into a tiny range.
            
            **Gradient Penalty (GP) Mechanics:**
            We solve this using the modern **Gradient Penalty** approach. We create a "random interpolate" image (a mix of real and fake pixels) and add a penalty to the loss if the gradient norm of the Critic at that point is anything other than 1. This "soft" constraint allows the weights to take any value they need to learn complex artistic features while still maintaining the mathematical stability required for the Earth Mover distance to be calculated correctly.
            """)

            st.markdown("### 🖼️ **IX. Weight Initialization & Kaiming Scaling**")
            st.markdown("""
            How you start the network matters as much as how you train it. If the initial weights are too large, the activations explode; if too small, they vanish.
            
            **Mean-Zero Gaussian Init:**
            We initialize our weights using a specialized Gaussian distribution with a mean of 0.0 and a standard deviation of 0.02. In the DCGAN research, this was found to be the "Goldilocks" zone for GAN stability. High-resolution art training is extremely sensitive to these initial conditions, as the adversarial game can be won or lost in the very first forward pass based on how much "signal" survives the journey through the deep convolutional layers.
            """)

            st.markdown("### ♾️ **X. Forward Pass Engineering**")
            st.markdown("""
            Behind the scenes, the training involves several advanced engineering choices:
            - **Adam Optimizer Momentum:** Standard Stochastic Gradient Descent (SGD) is too rigid for GANs. We use **Adam**, which maintains a "weighted moving average" of past gradients (momentum), allowing the optimization to "glide" over noisy local minima and handle the high-variance loss surfaces of GANs.
            - **Instance Normalization:** Unlike Batch Norm (which looks at many images at once), Instance Norm looks at each image individually. This has been technically shown to preserve stylistic features and avoid the "washed out" look found in many basic GANs, as it prevents the statistics of one image from "bleeding" into another during the training pass.
            - **Hyperparameter Stability:** We use fixed learning rates of 0.0002 and specific beta values ($\beta_1=0.5, \beta_2=0.999$) that have been empirically proven to provide the highest stability for Deep Convolutional GAN architectures.
            """)
            
            st.info("""
            🔬 **Memory Optimization**: This app uses `Linear` to `Conv2d` projection instead of heavy fully connected layers to stay within the **4GB VRAM** target while maintaining a high feature count in deep layers.
            """)

if __name__ == "__main__":
    # Initialize session state
    if 'trained' not in st.session_state:
        st.session_state['trained'] = False
    
    main()