Upload 2 files

generator.py

@@ -0,0 +1,360 @@
+import os
+import argparse
+import numpy as np
+import torch
+import torch.nn as nn
+import matplotlib.pyplot as plt
+from pathlib import Path
+from tkinter import *
+from PIL import ImageTk, Image
+import random
+from safetensors.torch import load_file
+
+# Generator model definition (must match the training architecture)
+class Generator(nn.Module):
+    def __init__(self, codings_size, image_size, image_channels):
+        super(Generator, self).__init__()
+        
+        self.fc = nn.Linear(codings_size, 6 * 6 * 256, bias=False)
+        self.bn1 = nn.BatchNorm1d(6 * 6 * 256)
+        self.leaky_relu = nn.LeakyReLU(0.2)
+        
+        self.conv_transpose1 = nn.ConvTranspose2d(256, 128, kernel_size=5, stride=1, padding=2, bias=False)
+        self.bn2 = nn.BatchNorm2d(128)
+        
+        self.conv_transpose2 = nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2, padding=2, output_padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(64)
+        
+        self.conv_transpose3 = nn.ConvTranspose2d(64, image_channels, kernel_size=5, stride=2, padding=2, output_padding=1, bias=False)
+        self.tanh = nn.Tanh()
+    
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.bn1(x)
+        x = self.leaky_relu(x)
+        x = x.view(-1, 256, 6, 6)
+        
+        x = self.conv_transpose1(x)
+        x = self.bn2(x)
+        x = self.leaky_relu(x)
+        
+        x = self.conv_transpose2(x)
+        x = self.bn3(x)
+        x = self.leaky_relu(x)
+        
+        x = self.conv_transpose3(x)
+        x = self.tanh(x)
+        
+        return x
+
+def load_model(model_path, device='cpu'):
+    """
+    Load the trained generator model from safetensors format.
+    
+    Args:
+        model_path: Path to the .safetensors model file
+        device: Device to load the model on ('cpu' or 'cuda')
+    
+    Returns:
+        Loaded generator model and configuration parameters
+    """
+    # Load state dict and metadata from safetensors
+    state_dict = load_file(model_path)
+    
+    # Load metadata from safetensors file
+    from safetensors import safe_open
+    with safe_open(model_path, framework="pt", device=str(device)) as f:
+        metadata = f.metadata()
+    
+    # Extract model configuration from metadata
+    codings_size = int(metadata['codings_size'])
+    image_size = int(metadata['image_size'])
+    image_channels = int(metadata['image_channels'])
+    
+    # Create generator model
+    model = Generator(codings_size, image_size, image_channels)
+    model.load_state_dict(state_dict)
+    model.to(device)
+    model.eval()
+    
+    print(f"Model configuration: codings_size={codings_size}, image_size={image_size}, image_channels={image_channels}")
+    
+    return model, codings_size, image_size, image_channels
+
+def generate_images(model, num_images, codings_size=100, seed=None, device='cpu'):
+    """
+    Generate images using the trained GAN generator model.
+    
+    Args:
+        model: Loaded PyTorch generator model
+        num_images: Number of images to generate
+        codings_size: Size of the latent vector (default: 100)
+        seed: Random seed for reproducibility
+        device: Device to run generation on
+    
+    Returns:
+        Generated images as numpy array (scaled to [0, 1])
+    """
+    if seed is not None:
+        torch.manual_seed(seed)
+        np.random.seed(seed)
+    
+    # Generate random noise as input
+    noise = torch.randn(num_images, codings_size, device=device)
+    
+    # Generate images
+    with torch.no_grad():
+        generated_images = model(noise)
+    
+    # Convert from CHW to HWC format and scale from [-1, 1] to [0, 1]
+    generated_images = generated_images.permute(0, 2, 3, 1).cpu().numpy()
+    generated_images = (generated_images + 1) / 2  # Scale to [0, 1]
+    
+    return generated_images
+
+def save_image_grid(images, output_path, grid_size=None):
+    """
+    Save generated images as a grid visualization.
+    
+    Args:
+        images: Array of generated images
+        output_path: Path to save the grid image
+        grid_size: Optional grid size (rows, cols). If None, auto-calculate square grid
+    """
+    num_images = images.shape[0]
+    
+    if grid_size is None:
+        # Auto-calculate square grid
+        grid_rows = int(np.sqrt(num_images))
+        grid_cols = int(np.ceil(num_images / grid_rows))
+    else:
+        grid_rows, grid_cols = grid_size
+    
+    fig = plt.figure(figsize=(grid_cols * 2, grid_rows * 2))
+    
+    for i in range(min(num_images, grid_rows * grid_cols)):
+        plt.subplot(grid_rows, grid_cols, i + 1)
+        
+        # Handle different image formats
+        if images.shape[-1] == 1:
+            # Grayscale
+            plt.imshow(images[i, :, :, 0], cmap='gray')
+        else:
+            # RGB or RGBA
+            plt.imshow(images[i])
+        
+        plt.axis('off')
+    
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches='tight')
+    plt.close()
+
+def save_individual_images(images, output_dir, prefix="generated"):
+    """
+    Save each generated image as a separate file.
+    
+    Args:
+        images: Array of generated images
+        output_dir: Directory to save individual images
+        prefix: Prefix for image filenames
+    """
+    output_dir = Path(output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    for i, img in enumerate(images):
+        # Convert to uint8 format (0-255)
+        img_uint8 = (img * 255).astype(np.uint8)
+        
+        # Save using matplotlib to handle RGBA correctly
+        output_path = output_dir / f"{prefix}_{i:04d}.png"
+        plt.imsave(output_path, img_uint8)
+    
+    print(f"Saved {len(images)} individual images to: {output_dir}")
+
+
+# ============ TKINTER UI MODE ============
+
+def run_gui(model_path, output_path):
+    """
+    Run Tkinter GUI for interactive image generation.
+    """
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Load model once at startup
+    print(f"Loading model from: {model_path}")
+    try:
+        model, codings_size, image_size, image_channels = load_model(model_path, device)
+        print("Model loaded successfully!")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        import traceback
+        traceback.print_exc()
+        return
+    
+    # Create output directory
+    output_dir = Path(output_path).parent
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Initialize Tkinter window
+    root = Tk()
+    root.title("CryptoPunk Generator")
+    root.columnconfigure([0, 1, 2, 3], minsize=200)
+    
+    # Create a placeholder image if output doesn't exist
+    if not os.path.exists(output_path):
+        fig = plt.figure(figsize=(4, 4))
+        plt.text(0.5, 0.5, 'Click a button to generate!', 
+                ha='center', va='center', fontsize=16)
+        plt.axis('off')
+        plt.savefig(output_path, dpi=150, bbox_inches='tight')
+        plt.close()
+    
+    # Load and display initial image
+    img = ImageTk.PhotoImage(Image.open(output_path))
+    panel = Label(root, image=img)
+    panel.grid(row=1, columnspan=4, sticky="nsew")
+    
+    def update_img():
+        """Update the displayed image"""
+        new_img = ImageTk.PhotoImage(Image.open(output_path))
+        panel.configure(image=new_img)
+        panel.image = new_img
+    
+    def generate(grid_size):
+        """Generate images in a grid"""
+        print(f"Generating {grid_size}x{grid_size} grid...")
+        n_img = grid_size * grid_size
+        seed = random.getrandbits(32)
+        
+        # Generate images
+        images = generate_images(model, n_img, codings_size, seed, device)
+        
+        # Create grid visualization
+        fig = plt.figure(figsize=(8, 8))
+        for i in range(n_img):
+            plt.subplot(grid_size, grid_size, i + 1)
+            plt.imshow(images[i, :, :, :])
+            plt.axis('off')
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=150, bbox_inches='tight')
+        plt.close()
+        
+        print(f"Generated with seed: {seed}")
+        update_img()
+    
+    # Create buttons
+    btn_1 = Button(root, text="Generate 1 cryptopunk", command=lambda: generate(1))
+    btn_3 = Button(root, text="Generate 3x3 cryptopunks", command=lambda: generate(3))
+    btn_5 = Button(root, text="Generate 5x5 cryptopunks", command=lambda: generate(5))
+    btn_q = Button(root, text="Terminate", command=root.quit)
+    
+    btn_1.grid(row=0, column=0, sticky="nsew")
+    btn_3.grid(row=0, column=1, sticky="nsew")
+    btn_5.grid(row=0, column=2, sticky="nsew")
+    btn_q.grid(row=0, column=3, sticky="nsew")
+    
+    print("\nGUI started! Click buttons to generate images.")
+    root.mainloop()
+
+
+# ============ CLI MODE ============
+
+def run_cli(args):
+    """
+    Run command-line interface for batch image generation.
+    """
+    # Check if model exists
+    if not os.path.exists(args.model_path):
+        print(f"Error: Model not found at {args.model_path}")
+        print("Please train the model first using trainer.py")
+        return
+    
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+    
+    # Load the trained model
+    print(f"Loading model from: {args.model_path}")
+    try:
+        model, codings_size, image_size, image_channels = load_model(args.model_path, device)
+        print("Model loaded successfully!")
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        import traceback
+        traceback.print_exc()
+        return
+    
+    # Calculate actual number of images for grid
+    if args.grid_size is not None:
+        num_images = args.grid_size * args.grid_size
+        grid_size = (args.grid_size, args.grid_size)
+        print(f"Generating {num_images} images in a {args.grid_size}x{args.grid_size} grid")
+    else:
+        num_images = args.num_images
+        grid_size = None
+        print(f"Generating {num_images} images")
+    
+    # Generate images
+    print("Generating images...")
+    images = generate_images(model, num_images, codings_size, args.seed, device)
+    print(f"Generated images shape: {images.shape}")
+    print(f"Image value range: [{images.min():.3f}, {images.max():.3f}]")
+    
+    # Create output directory if needed
+    output_dir = Path(args.output_path).parent
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Save grid visualization
+    save_image_grid(images, args.output_path, grid_size)
+    print(f"Grid image saved to: {args.output_path}")
+    
+    # Optionally save individual images
+    if args.save_individual:
+        save_individual_images(images, args.individual_output_dir)
+    
+    print("\nGeneration complete!")
+    if args.seed is not None:
+        print(f"Seed used: {args.seed} (use same seed to reproduce these images)")
+
+
+# ============ MAIN ============
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate images using trained GAN model")
+    parser.add_argument("--gui", action="store_true",
+                        help="Launch Tkinter GUI interface (default if no other args)")
+    parser.add_argument("--model_path", type=str, default="./models/generator_model.safetensors",
+                        help="Path to the trained generator model (.safetensors file)")
+    parser.add_argument("--output_path", type=str, default="./generated/output.png",
+                        help="Path to save the generated image grid")
+    parser.add_argument("--num_images", type=int, default=16,
+                        help="Number of images to generate (CLI mode, default: 16)")
+    parser.add_argument("--grid_size", type=int, default=None,
+                        help="Grid size N for NxN layout (CLI mode)")
+    parser.add_argument("--seed", type=int, default=None,
+                        help="Random seed for reproducibility (CLI mode only)")
+    parser.add_argument("--save_individual", action="store_true",
+                        help="Save each generated image as a separate file (CLI mode)")
+    parser.add_argument("--individual_output_dir", type=str, default="./generated/individual/",
+                        help="Directory to save individual images (CLI mode)")
+    
+    args = parser.parse_args()
+    
+    # Determine mode: GUI if --gui flag or if no CLI-specific args provided
+    cli_args_provided = (args.grid_size is not None or 
+                         args.num_images != 16 or 
+                         args.seed is not None or 
+                         args.save_individual)
+    
+    if args.gui or not cli_args_provided:
+        # GUI mode
+        run_gui(args.model_path, args.output_path)
+    else:
+        # CLI mode
+        run_cli(args)
+
+
+if __name__ == "__main__":
+    main()

trainer.py

@@ -0,0 +1,244 @@
+import os
+import glob
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms
+from PIL import Image
+import matplotlib.pyplot as plt
+from safetensors.torch import save_file
+
+def plot_multiple_images(images, n_cols, epoch):
+    n_cols = n_cols or len(images)
+    n_rows = (len(images) - 1) // n_cols + 1
+    # Convert from CHW to HWC format for plotting
+    images = images.permute(0, 2, 3, 1).cpu().numpy()
+    if images.shape[-1] == 1:
+        images = np.squeeze(images, axis=-1)
+    plt.figure(figsize=(n_cols, n_rows))
+    for index, image in enumerate(images):
+        image = ((image + 1) / 2) # scale back
+        plt.subplot(n_rows, n_cols, index + 1)
+        plt.imshow(image, cmap="binary")
+        plt.axis("off")
+    plt.savefig(f'{args.images_output_path}epoch_{epoch}.png')
+    plt.close()  # Close the figure to free memory
+
+class ImageDataset(Dataset):
+    def __init__(self, file_paths, image_size, image_channels):
+        self.file_paths = file_paths
+        self.image_size = image_size
+        self.image_channels = image_channels
+        self.transform = transforms.Compose([
+            transforms.Resize((image_size, image_size)),
+            transforms.ToTensor(),
+            transforms.Normalize([0.5] * image_channels, [0.5] * image_channels)  # Scale to [-1, 1]
+        ])
+    
+    def __len__(self):
+        return len(self.file_paths)
+    
+    def __getitem__(self, idx):
+        img_path = self.file_paths[idx]
+        image = Image.open(img_path).convert('RGBA' if self.image_channels == 4 else 'RGB')
+        image = self.transform(image)
+        return image
+
+def get_dataloader(inputs, batch_size, image_size, image_channels):
+    if type(inputs) == dict:
+        file_paths = inputs["paths"].tolist()
+    else:
+        file_paths = glob.glob(f"{inputs}/*")
+    
+    dataset = ImageDataset(file_paths, image_size, image_channels)
+    dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, drop_last=True, num_workers=2)
+    return dataloader
+
+def discriminator_loss(real_output, fake_output, criterion):
+    real_loss = criterion(real_output, torch.ones_like(real_output))
+    fake_loss = criterion(fake_output, torch.zeros_like(fake_output))
+    total_loss = real_loss + fake_loss
+    return total_loss
+
+def generator_loss(fake_output, criterion):
+    return criterion(fake_output, torch.ones_like(fake_output))
+
+def train_step(images, batch_size, codings_size, generator, discriminator, gen_optimizer, disc_optimizer, criterion, device):
+    noise = torch.randn(batch_size, codings_size, device=device)
+    
+    # Train Discriminator
+    disc_optimizer.zero_grad()
+    generated_images = generator(noise)
+    real_output = discriminator(images)
+    fake_output = discriminator(generated_images.detach())
+    disc_loss = discriminator_loss(real_output, fake_output, criterion)
+    disc_loss.backward()
+    disc_optimizer.step()
+    
+    # Train Generator
+    gen_optimizer.zero_grad()
+    fake_output = discriminator(generated_images)
+    gen_loss = generator_loss(fake_output, criterion)
+    gen_loss.backward()
+    gen_optimizer.step()
+    
+    return gen_loss.item(), disc_loss.item()
+
+def train(dataloader, epochs, batch_size, codings_size, generator, discriminator, gen_optimizer, disc_optimizer, criterion, device):
+    generator.train()
+    discriminator.train()
+    
+    for epoch in range(epochs):
+        for image_batch in dataloader:
+            image_batch = image_batch.to(device)
+            gen_loss, disc_loss = train_step(image_batch, batch_size, codings_size, generator, discriminator, 
+                                             gen_optimizer, disc_optimizer, criterion, device)
+        
+        print(f"Epoch {epoch+1}/{epochs} - Gen Loss: {gen_loss:.4f}, Disc Loss: {disc_loss:.4f}")
+        if args.images_output_path:
+            generator.eval()
+            with torch.no_grad():
+                noise = torch.randn(batch_size, codings_size, device=device)
+                display_images = generator(noise)
+                plot_multiple_images(display_images, 8, epoch)
+            generator.train()
+
+class Generator(nn.Module):
+    def __init__(self, codings_size, image_size, image_channels):
+        super(Generator, self).__init__()
+        
+        self.fc = nn.Linear(codings_size, 6 * 6 * 256, bias=False)
+        self.bn1 = nn.BatchNorm1d(6 * 6 * 256)
+        self.leaky_relu = nn.LeakyReLU(0.2)
+        
+        self.conv_transpose1 = nn.ConvTranspose2d(256, 128, kernel_size=5, stride=1, padding=2, bias=False)
+        self.bn2 = nn.BatchNorm2d(128)
+        
+        self.conv_transpose2 = nn.ConvTranspose2d(128, 64, kernel_size=5, stride=2, padding=2, output_padding=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(64)
+        
+        self.conv_transpose3 = nn.ConvTranspose2d(64, image_channels, kernel_size=5, stride=2, padding=2, output_padding=1, bias=False)
+        self.tanh = nn.Tanh()
+    
+    def forward(self, x):
+        x = self.fc(x)
+        x = self.bn1(x)
+        x = self.leaky_relu(x)
+        x = x.view(-1, 256, 6, 6)
+        
+        x = self.conv_transpose1(x)
+        x = self.bn2(x)
+        x = self.leaky_relu(x)
+        
+        x = self.conv_transpose2(x)
+        x = self.bn3(x)
+        x = self.leaky_relu(x)
+        
+        x = self.conv_transpose3(x)
+        x = self.tanh(x)
+        
+        return x
+
+class Discriminator(nn.Module):
+    def __init__(self, image_size, image_channels):
+        super(Discriminator, self).__init__()
+        
+        self.conv1 = nn.Conv2d(image_channels, 64, kernel_size=4, stride=2, padding=1)
+        self.leaky_relu1 = nn.LeakyReLU(0.2)
+        self.dropout1 = nn.Dropout(0.4)
+        
+        self.conv2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1)
+        self.leaky_relu2 = nn.LeakyReLU(0.2)
+        self.dropout2 = nn.Dropout(0.4)
+        
+        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
+        self.leaky_relu3 = nn.LeakyReLU(0.2)
+        self.dropout3 = nn.Dropout(0.4)
+        
+        self.global_avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(256, 1)
+        self.sigmoid = nn.Sigmoid()
+    
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.leaky_relu1(x)
+        x = self.dropout1(x)
+        
+        x = self.conv2(x)
+        x = self.leaky_relu2(x)
+        x = self.dropout2(x)
+        
+        x = self.conv3(x)
+        x = self.leaky_relu3(x)
+        x = self.dropout3(x)
+        
+        x = self.global_avg_pool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        x = self.sigmoid(x)
+        
+        return x
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--data_path", default="./data/attributes.csv", help="Path to dataset (attributes.csv)")
+    parser.add_argument("--images_path", default="./data/images/", help="Path to images")
+    parser.add_argument("--model_output_path", default="./models/", help="Path to output the generator model")
+    parser.add_argument("--images_output_path", default="./gen_images/", help="Path to output generated images during training")
+    parser.add_argument("--codings_size", type=int, default=100, help="Size of the latent z vector")
+    parser.add_argument("--image_size", type=int, default=24, help="Images size")
+    parser.add_argument("--image_channels", type=int, default=4, help="Images channels")
+    parser.add_argument("--batch_size", type=int, default=16, help="Input batch size")
+    parser.add_argument("--epochs", type=int, default=50, help="Number of epochs")
+    args = parser.parse_args()
+    print(args)
+
+    # Set device
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Using device: {device}")
+
+    if args.images_output_path and (os.path.exists(args.images_output_path) == False):
+        print(f"Saving generated images during training at: {args.images_output_path}")
+        os.mkdir(args.images_output_path)
+
+    print("Loading the dataset...")
+    df = pd.read_csv(args.data_path)
+    df.id = df.id.apply(lambda x: f"{args.images_path}punk{x:03d}.png")
+
+    print("Creating PyTorch DataLoader...")
+    dataloader = get_dataloader({"paths": df.id}, args.batch_size, args.image_size, args.image_channels)
+    
+    generator = Generator(args.codings_size, args.image_size, args.image_channels).to(device)
+    print("Generator architecture:")
+    print(generator)
+
+    discriminator = Discriminator(args.image_size, args.image_channels).to(device)
+    print("Discriminator architecture:")
+    print(discriminator)
+
+    gen_optimizer = optim.RMSprop(generator.parameters(), lr=0.001)
+    disc_optimizer = optim.RMSprop(discriminator.parameters(), lr=0.001)
+    criterion = nn.BCELoss()
+
+    print("Training model...")
+    train(dataloader, args.epochs, args.batch_size, args.codings_size, generator, discriminator, 
+          gen_optimizer, disc_optimizer, criterion, device)
+
+    print(f"Saving model at: {args.model_output_path}...")
+    os.makedirs(args.model_output_path, exist_ok=True)
+    model_path = args.model_output_path if args.model_output_path.endswith('.safetensors') else os.path.join(args.model_output_path, 'generator_model.safetensors')
+    
+    # Save the generator model in safetensors format
+    # Metadata is stored as strings in safetensors
+    metadata = {
+        'codings_size': str(args.codings_size),
+        'image_size': str(args.image_size),
+        'image_channels': str(args.image_channels)
+    }
+    save_file(generator.state_dict(), model_path, metadata=metadata)
+    print(f"Model saved to: {model_path}")