You can deploy this model and start earning money today!

egpivo/
biggan-mnist-finetuned

Name: biggan-mnist-finetuned
Availability: InStock

biggan

mit

Views

No views yet

BigGAN Fine-Tuned on MNIST

This model is a fine-tuned version of BigGAN (biggan-deep-256) trained on the MNIST dataset for generating realistic digit images. It supports class-conditioned generation for digits 0 through 9.

Model Details

Base Model: Pre-trained BigGAN (biggan-deep-256)
Dataset: MNIST (10 classes, grayscale images)
Objective: Generate high-quality, class-conditioned digit images.

How to Use

Installation

Install the required dependencies:

pip install torch torchvision pytorch-pretrained-biggan

Generate Images

python

import torch
from huggingface_hub import hf_hub_download
from pytorch_pretrained_biggan import BigGAN, truncated_noise_sample
from torchvision.utils import make_grid
import matplotlib.pyplot as plt

model = BigGAN.from_pretrained("biggan-deep-256")

checkpoint_path = hf_hub_download(
    repo_id="egpivo/biggan-mnist-finetuned",
    filename="biggan_mnist_embedding.pth"
)
state_dict = torch.load(checkpoint_path, map_location="cpu")

mnist_embedding = torch.nn.Embedding(10, 1000)
mnist_embedding.load_state_dict({"weight": state_dict["weight"]})

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
model.eval()
mnist_embedding.to(device)

batch_size = 16
truncation = 0.7
noise = torch.randn(batch_size, 128, device=device)

labels = torch.randint(0, 10, (batch_size,), device=device, dtype=torch.long)
print("Generated labels:", labels.cpu().numpy())

class_embedding = mnist_embedding(labels)  # Use separate embedding layer

with torch.no_grad():
    generated_images = model(noise, class_embedding, truncation)

generated_images = (generated_images + 1) / 2
grid = make_grid(generated_images, nrow=4)
plt.imshow(grid.permute(1, 2, 0).cpu())
plt.axis("off")
plt.show()

Acknowledgments

Fine-tuned using the MNIST dataset. Special thanks to the creators of BigGAN and the MNIST dataset.