ohnah/verity-fastgan-image-gen

Verity FastGAN

A FastGAN generator trained on a small custom image dataset of "Verity" (the scary yellow smiley guy), using the bingchenlll/FastGAN-pytorch implementation of "Towards Faster and Stabilized GAN Training for High-fidelity Few-shot Image Synthesis" (ICLR 2021).

Generates 1024x1024 images from random noise. Fast enough to run on CPU.

Model details

• Architecture: FastGAN generator (skip-layer excitation, few-shot GAN)
• Resolution: 1024x1024
• Training data: a small custom dataset (see Training data below)
• Framework: PyTorch

How to use

This model's weights are only usable together with the Generator class defined in models.py from the FastGAN-pytorch repo — the checkpoint is a plain state_dict, not a standalone package, so you need that file alongside your inference code.

1. Get the code

Clone the FastGAN-pytorch repo (for models.py) and download this model's checkpoint file:

git clone https://github.com/bingchenlll/FastGAN-pytorch.git
cd FastGAN-pytorch
# place the downloaded .pth checkpoint from this repo somewhere inside, e.g. ./verity.pth

2. Install dependencies

pip install torch torchvision

No other dependencies are needed for inference (training-only packages like lpips are not required).

3. Run inference

Save the script below as inference.py in the same folder as models.py, then run:

python inference.py --ckpt verity.pth --out result.png

Pass --seed for a reproducible image, or omit it for a random one each run:

python inference.py --ckpt verity.pth --out result.png --seed 1234

inference.py

#!/usr/bin/env python3
"""
inference.py

Standalone, CPU-only inference script for a FastGAN generator trained with
bingchenlll/FastGAN-pytorch (https://github.com/bingchenlll/FastGAN-pytorch).

Generates a single image per run. By default uses a random seed; pass --seed
for a reproducible one.

Usage:
    python inference.py --ckpt path/to/checkpoint.pth --out out.png
    python inference.py --ckpt path/to/checkpoint.pth --out out.png --seed 1234
"""

import argparse
import os
import random
import sys

import torch
from torchvision.utils import save_image

os.environ["CUDA_VISIBLE_DEVICES"] = ""
DEVICE = torch.device("cpu")

try:
    from models import Generator
except ImportError:
    sys.exit(
        "Could not import 'Generator' from models.py.\n"
        "Run this script from inside your FastGAN-pytorch repo directory, "
        "or add that directory to PYTHONPATH."
    )

CKPT_KEY = "g"


def parse_args():
    parser = argparse.ArgumentParser(description="Generate one image from a trained FastGAN generator (CPU-only).")
    parser.add_argument("--ckpt", type=str, required=True, help="Path to the generator checkpoint (.pth) file.")
    parser.add_argument("--out", type=str, default="output.png", help="Output image path.")
    parser.add_argument("--seed", type=int, default=None, help="Random seed. Omit for a random seed each run.")
    parser.add_argument("--im_size", type=int, default=1024, help="Generator output resolution.")
    parser.add_argument("--nz", type=int, default=256, help="Latent (noise) vector dimensionality.")
    parser.add_argument("--ngf", type=int, default=64, help="Generator base feature width.")
    parser.add_argument("--nc", type=int, default=3, help="Number of image channels (3 for RGB).")
    return parser.parse_args()


def set_seed(seed: int):
    random.seed(seed)
    torch.manual_seed(seed)


def load_generator(ckpt_path: str, im_size: int, nz: int, ngf: int, nc: int) -> torch.nn.Module:
    generator = Generator(ngf=ngf, nz=nz, nc=nc, im_size=im_size)
    generator.to(DEVICE)

    checkpoint = torch.load(ckpt_path, map_location=DEVICE)

    if isinstance(checkpoint, dict) and CKPT_KEY in checkpoint:
        state_dict = checkpoint[CKPT_KEY]
    elif isinstance(checkpoint, dict) and all(
        isinstance(v, torch.Tensor) for v in checkpoint.values()
    ):
        state_dict = checkpoint
    else:
        sys.exit(
            f"Could not find generator weights under key '{CKPT_KEY}' in the "
            f"checkpoint, and it doesn't look like a plain state_dict either.\n"
            f"Checkpoint top-level keys: {list(checkpoint.keys()) if isinstance(checkpoint, dict) else type(checkpoint)}"
        )

    if any(k.startswith("module.") for k in state_dict.keys()):
        state_dict = {k.replace("module.", "", 1): v for k, v in state_dict.items()}

    generator.load_state_dict(state_dict)
    generator.eval()
    return generator


def denormalize(img: torch.Tensor) -> torch.Tensor:
    # Matches train.py's own preview-image convention: fixed linear remap
    # from tanh's [-1, 1] range to [0, 1].
    return img.add(1).mul(0.5)


def main():
    args = parse_args()

    seed = args.seed if args.seed is not None else random.randint(0, 2**31 - 1)
    set_seed(seed)
    print(f"Using seed: {seed}")

    generator = load_generator(
        ckpt_path=args.ckpt,
        im_size=args.im_size,
        nz=args.nz,
        ngf=args.ngf,
        nc=args.nc,
    )

    noise = torch.randn(1, args.nz, device=DEVICE)

    with torch.no_grad():
        result = generator(noise)

    image = result[0] if isinstance(result, (tuple, list)) else result
    image = denormalize(image)
    save_image(image, args.out, nrow=1)
    print(f"Saved image to: {args.out}")


if __name__ == "__main__":
    main()

Checkpoint format

The checkpoint is a dict saved by the repo's train.py as:

torch.save({'g': netG.state_dict(), 'd': netD.state_dict(), 'g_ema': avg_param_G, ...}, path)

This model's published .pth file uses the 'g' key (the raw generator weights, not the EMA-averaged ones).

Training data

The model was trained on random pictures of Verity (again, the smiley yellow ball from MINECRAFT!!) from Horror Skunx's "It's me, It's Verity" music video.

Limitations and bias

• Trained on a very small dataset, following the few-shot GAN setup described in the FastGAN paper. Output diversity and quality are bounded by how much the training set covered.
• Not evaluated for demographic or content bias; use judgment before using outputs beyond experimentation.
• Research/personal use; see the base repo's own terms before commercial use.

Why?

I made this for fun, because I was bored. This is also the first FastGAN model I trained. Anyway, I hope you have fun with this horrible model :)

Citation

@inproceedings{liu2021towards,
  title={Towards Faster and Stabilized {GAN} Training for High-fidelity Few-shot Image Synthesis},
  author={Bingchen Liu and Yizhe Zhu and Kunpeng Song and Ahmed Elgammal},
  booktitle={International Conference on Learning Representations},
  year={2021}
}