generate_examples.py

import argparse
import os

import matplotlib.pyplot as plt
import torch
from misc import generate_text, generate_text_conditioned, get_datasets_and_loaders
from vqvae.transformer import GPT2, ConditionnedGPT2
from vqvae.vqvae import VectorQuantizedVAE

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="VQ-VAE training and visualization script"
    )
    # Data Handling arguments
    parser.add_argument(
        "--dataset",
        type=str,
        default="mnist",
        choices=["mnist", "fashion_mnist", "cifar10"],
        help="Dataset to use",
    )
    parser.add_argument(
        "--data-root", type=str, default=".", help="Root directory for datasets"
    )
    parser.add_argument(
        "--batch-size", type=int, default=128, help="Input batch size for training"
    )
    parser.add_argument(
        "--num-workers", type=int, default=2, help="Number of workers for data loading"
    )

    # VQ-VAE Model arguments
    parser.add_argument(
        "--embedding-dim",
        type=int,
        default=64,
        help="Dimensionality of VQ embedding vectors",
    )
    parser.add_argument(
        "--num-embeddings",
        type=int,
        default=128,
        help="Number of VQ embedding vectors (codebook size K)",
    )

    parser.add_argument(
        "--hidden-dims",
        type=int,
        default=128,
        help="Hidden dimensions in Encoder/Decoder CNNs",
    )

    # Transformer Model arguments
    parser.add_argument(
        "--block-size",
        type=int,
        default=256,
        help="Block size for transformer",
    )

    parser.add_argument(
        "--n-heads",
        type=int,
        default=8,
        help="Number of attention heads in transformer",
    )
    parser.add_argument(
        "--n-layers",
        type=int,
        default=6,
        help="Number of transformer layers",
    )
    parser.add_argument(
        "--gpt-embedding-dim",
        type=int,
        default=128,
        help="Dimensionality of VQ embedding vectors",
    )
    parser.add_argument(
        "--dropout",
        type=float,
        default=0.1,
        help="Dropout rate for transformer",
    )
    parser.add_argument("--conditioned", action="store_true")
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda", "cpu"],
        help="Device to use (cuda or cpu)",
    )

    parser.add_argument(
        "--trained-vqvae-path",
        type=str,
        required=True,
        help="Trained VQ-VAE path",
    )

    parser.add_argument(
        "--trained-prior-path",
        type=str,
        required=True,
        help="Trained prior path",
    )
    args = parser.parse_args()

    device = torch.device(
        args.device if torch.cuda.is_available() and args.device == "cuda" else "cpu"
    )

    train_loader, _, n_channels, img_size = get_datasets_and_loaders(
        args.dataset, args.data_root, args.batch_size, args.num_workers
    )
    print(args.conditioned)
    model = VectorQuantizedVAE(
        n_channels,
        args.embedding_dim,
        args.num_embeddings,
        1.0,
        args.hidden_dims,
    ).to(device)

    trained_vqvae_path = os.path.join(args.dataset, args.trained_vqvae_path)

    model.load_state_dict(
        torch.load(trained_vqvae_path + "/best_vqvae.pt", map_location=device)
    )
    model = model.to(device)
    model.eval()

    # Determine latent dimensions
    dummy_input = torch.randn(1, n_channels, img_size, img_size).to(device)
    with torch.no_grad():
        encoder_output_shape = model.encoder(dummy_input).shape
    latent_h_vq, latent_w_vq = encoder_output_shape[2], encoder_output_shape[3]
    print(f"Latent grid dimensions after encoder: {latent_h_vq}x{latent_w_vq}")

    if args.conditioned is False:
        n_classes = None
        prior = GPT2(
            args.num_embeddings,
            args.block_size,
            args.gpt_embedding_dim,
            args.n_heads,
            args.n_layers,
            args.dropout,
            True,
        )
    else:
        dataset_obj = train_loader.dataset
        n_classes = len(dataset_obj.classes)  # type: ignore

        prior = ConditionnedGPT2(
            args.num_embeddings,
            n_classes,
            args.block_size,
            args.gpt_embedding_dim,
            args.gpt_embedding_dim,
            args.n_heads,
            args.n_layers,
            args.dropout,
            True,
        )

    trained_prior_path = os.path.join(args.dataset, args.trained_prior_path)
    prior.load_state_dict(
        torch.load(trained_prior_path + "/best_prior.pt", map_location=device)
    )
    prior = prior.to(device)
    prior.eval()

    random_ids = torch.randint(
        0,
        model.quantizer.num_embeddings,
        (16, 1),
        device=device,
    )

    if args.conditioned:
        assert n_classes is not None
        random_class_ids = torch.randint(
            0,
            n_classes,
            (16,),
            device=device,
        )
        generated_indices = generate_text_conditioned(
            prior,
            random_ids,
            random_class_ids,
            generation_length=latent_h_vq * latent_w_vq - 1,
            temperature=1.0,
            top_k=None,
        )
    else:
        generated_indices = generate_text(
            prior,
            random_ids,
            latent_h_vq * latent_w_vq - 1,
            temperature=0.7,
            top_k=None,
        )

    generated_indices = generated_indices.reshape(-1, latent_h_vq, latent_w_vq)
    generated_images = model.reconstruct_from_indices(generated_indices).squeeze(1)

    cmap_gen = "gray" if n_channels == 1 else None
    for i in range(10):
        plt.subplot(1, 10, i + 1)
        plt.imshow(generated_images[i].cpu().numpy(), cmap=cmap_gen)
        if args.conditioned:
            plt.title(f"Class: {random_class_ids[i].item()}")  # type: ignore
        plt.axis("off")
    plt.tight_layout()
    plt.show()