mythcell/galaxy-vqvae

Galaxy VQ-VAE

A minimalist Gradio interface for generating synthetic galaxy images using a VQ-VAE + Transformer Prior model trained on the Galaxy10 DECaLS dataset. The generator produces one image for each of the 10 classes in the dataset.

Architecture

The original model consists of two primary components:

1. 1.VQ-VAE: All-convolutional encoder and decoder for compressing $64 \times 64$ RGB images into a discrete latent space of codebook indices. The latent space is a $8 \times 8$ grid of 128-dimensional codebook vectors.
2. 2.Transformer Prior: An autoregressive model that learns the probability distribution of these discrete codes, conditioned on galaxy morphological classes. The prior is trained with class dropout to allow for CFG-style inference.

The inference model galaxy_generator.pt uses only the VQ-VAE decoder and the (frozen) VQ codebook. We first use the Transformer prior to autoregressive generate a sequence of codebook indices. We then map these to codebook vectors via the trained VQ codebook, then finally pass the codebook vectors through the VQ-VAE decoder to output the final image.

Parameters

Conditioning

• —CFG Scale (Classifier-Free Guidance): Controls the strength of the class conditioning. Higher values force the model to produce images that more strongly represent the features of the selected class, but may reduce overall image diversity or introduce artifacts. Default value is 3.0.

Temperature (Linear Annealing)

The model uses linear temperature annealing during the autoregressive generation of the latent code sequence.

• —Temperature Start: The initial sampling temperature. Higher values increase randomness at the beginning of the codebook sequence. Default value is 1.0.
• —Temperature End: The final sampling temperature. Lower values (e.g., < 1.0) help the model settle on a clear structure as sequence of codes nears completion. Default value is 0.7.

Sampling

• —Top-K: Limits the sampling pool to the top $K$ most likely next tokens. A value of $0$ disables this filter.
• —Top-P (Nucleus Sampling): Limits the sampling pool to the smallest set of tokens whose cumulative probability exceeds $P$. This allows for more dynamic sampling than fixed Top-K.

Generation Settings

• —Seed: A fixed integer to ensure reproducibility. Set to -1 to generate a random seed.

Install

Simply install with pip or uv, ideally in a virtual environment:

pip install -r requirements.txt
python app.py