metadata
license: mit
This model was trained using the 8888 images of the Anky Genesis NFT Collection, and its mission is to transform an image into pixel art, like so:
The code used for training it is the following:
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms
from PIL import Image
import numpy as np
# Custom dataset for loading the images
class PixelArtDataset(Dataset):
def __init__(self, image_folder, transform=None):
self.image_folder = image_folder
self.transform = transform
self.image_files = [f"{i}.png" for i in range(1, 8889)]
# Debug: Check if images are correctly listed
print(f"Total images found: {len(self.image_files)}")
def __len__(self):
return len(self.image_files)
def __getitem__(self, idx):
img_path = os.path.join(self.image_folder, self.image_files[idx])
image = Image.open(img_path).convert("RGB")
if self.transform:
image = self.transform(image)
return image, image
# Define the neural network
class PixelArtGenerator(nn.Module):
def __init__(self):
super(PixelArtGenerator, self).__init__()
print("Initializing PixelArtGenerator Model...")
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
def train(model, dataloader, criterion, optimizer, device, epochs=50):
print("Starting training...")
model.train()
for epoch in range(epochs):
running_loss = 0.0
print(f"Epoch [{epoch+1}/{epochs}] starting...")
for batch_idx, (input_images, target_images) in enumerate(dataloader):
input_images, target_images = input_images.to(device), target_images.to(device)
optimizer.zero_grad()
outputs = model(input_images)
loss = criterion(outputs, target_images)
loss.backward()
optimizer.step()
running_loss += loss.item()
# Debug: Print progress for every batch
if batch_idx % 10 == 0:
print(f"Epoch [{epoch+1}/{epochs}], Batch [{batch_idx+1}/{len(dataloader)}], Loss: {loss.item():.4f}")
print(f"Epoch [{epoch+1}/{epochs}] completed with Loss: {running_loss/len(dataloader):.4f}")
def create_pixel_art(model, input_image_path, output_image_path, device):
print("Creating pixel art...")
model.eval()
transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
image = Image.open(input_image_path).convert("RGB")
input_image = transform(image).unsqueeze(0).to(device)
with torch.no_grad():
output_image = model(input_image).squeeze(0).cpu().numpy()
output_image = np.transpose(output_image, (1, 2, 0))
output_image = (output_image * 0.5 + 0.5) * 255.0
output_image = np.clip(output_image, 0, 255).astype(np.uint8)
output_image = Image.fromarray(output_image)
output_image.save(output_image_path)
print(f"Pixel art saved to {output_image_path}")
if __name__ == "__main__":
# Transform for input images
print("Setting up image transformations...")
transform = transforms.Compose([
transforms.Resize((64, 64)), # Resize to 64x64 for input
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
# Load dataset
print("Loading dataset...")
image_folder = "./" # Change this to your images folder path
dataset = PixelArtDataset(image_folder, transform)
dataloader = DataLoader(dataset, batch_size=8, shuffle=True) # Reduce batch size for debugging
# Check for GPU availability
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")
# Initialize the model, criterion, and optimizer
model = PixelArtGenerator().to(device)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.0002)
# Enable data parallelism if multiple GPUs are available
if torch.cuda.device_count() > 1:
print(f"Using {torch.cuda.device_count()} GPUs")
model = nn.DataParallel(model)
# Train the model
train(model, dataloader, criterion, optimizer, device, epochs=50)
# Save the model
torch.save(model.state_dict(), "pixel_art_generator.pth")
print("Model saved as 'pixel_art_generator.pth'")
# Create pixel art from a new input image
input_image_path = "input_image.png" # Path to the high-resolution input image
output_image_path = "pixel_art.png" # Path to save the generated pixel art
create_pixel_art(model, input_image_path, output_image_path, device)
print("Pixel art creation completed.")
The training happened on a Cognition PRO called poiesis. It consisted of 50 epochs, and it lasted for about 4 hours running on 2x NVIDIA RTX 4090.
Its intended usage is for it to transform any image into its corresponding in pixels, as you can see on this one.
For running it like such, you can run the following python code on the containing folder of the model (for transforming an image called pfp.png):
import torch
import torch.nn as nn
from PIL import Image
import numpy as np
from torchvision import transforms
import os
# Define the neural network (same as the one used during training)
class PixelArtGenerator(nn.Module):
def __init__(self):
super(PixelArtGenerator, self).__init__()
print("Initializing PixelArtGenerator Model...")
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(256),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(),
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
nn.BatchNorm2d(64),
nn.ReLU(),
nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
def forward(self, x):
x = self.encoder(x)
x = self.decoder(x)
return x
def create_pixel_art(model, input_image_path, output_image_path, device):
print(f"Creating pixel art for {input_image_path}...")
# Check if the input image file exists
if not os.path.isfile(input_image_path):
print(f"Error: Input image file '{input_image_path}' not found.")
return
model.eval()
print("Model set to evaluation mode.")
# Define the transformation for the input image
transform = transforms.Compose([
transforms.Resize((64, 64)),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])
print("Image transformation defined.")
# Load and preprocess the input image
image = Image.open(input_image_path).convert("RGB")
input_image = transform(image).unsqueeze(0).to(device)
print(f"Input image '{input_image_path}' loaded and preprocessed.")
# Generate pixel art using the model
with torch.no_grad():
output_image = model(input_image).squeeze(0).cpu().numpy()
print("Pixel art generated by the model.")
# Post-process and save the output image
output_image = np.transpose(output_image, (1, 2, 0))
output_image = (output_image * 0.5 + 0.5) * 255.0
output_image = np.clip(output_image, 0, 255).astype(np.uint8)
output_image = Image.fromarray(output_image)
# Scale up the image to iPhone 11 width (828 pixels)
scaled_output_image = output_image.resize((828, int(828 * output_image.size[1] / output_image.size[0])), Image.NEAREST)
scaled_output_image.save(output_image_path)
print(f"Pixel art saved to '{output_image_path}'.")
if __name__ == "__main__":
print("Starting pixel art generation script...")
# Load the trained model
model = PixelArtGenerator()
model_path = "pixel_art_generator.pth" # Path to the saved model
print(f"Loading model from '{model_path}'...")
# Load model with handling for DataParallel
state_dict = torch.load(model_path)
if 'module.' in list(state_dict.keys())[0]:
# Remove 'module.' prefix if model was saved with DataParallel
state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
model.load_state_dict(state_dict)
print("Model loaded successfully.")
# Check for GPU availability
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
print(f"Using device: {device}")
# Define the input and output paths for the single image
input_image_path = "pfp.jpeg" # Path to the input image
output_image_path = "pfp_pixelated.png" # Path to save the generated pixel art
# Create pixel art for the single image
create_pixel_art(model, input_image_path, output_image_path, device)
print("Pixel art creation completed for the single image.")
Hope you enjoy, and any questions that you may have, feel free to reach out to @jpfraneto on telegram.
If you want to contribute to Anky, we have plenty of compute available, and a powerful story (and intention) that puts the unfolding of AI at the core of our experience as humans.
Think of it as a playground for your inner child, with boundless potential.
Our farcaster channel is here: https://warpcast.com/~/channel/anky
Your uniqueness is a gift.
🎩