GitHub - huggingface/open-muse: Open reproduction of MUSE for fast text2image generation. (original) (raw)

An open-reproduction effort to reproduce the transformer based MUSE model for fast text2image generation.

Demo

👉 https://huggingface.co/spaces/openMUSE/MUSE

Goal

This repo is for reproduction of the MUSE model. The goal is to create a simple and scalable repo, to reproduce MUSE and build knowedge about VQ + transformers at scale. We will use deduped LAION-2B + COYO-700M dataset for training.

Project stages:

  1. Setup the codebase and train a class-conditional model on imagenet.
  2. Conduct text2image experiments on CC12M.
  3. Train improved VQGANs models.
  4. Train the full (base-256) model on LAION + COYO.
  5. Train the full (base-512) model on LAION + COYO.

All the artifacts of this project will be uploaded to the openMUSE organization on the huggingface hub.

Usage

Installation

First create a virtual environment and install the repo using:

git clone https://github.com/huggingface/muse cd muse pip install -e ".[extra]"

You'll need to install PyTorch and torchvision manually. We are using torch==1.13.1 with CUDA11.7 for training.

For distributed data parallel training we use accelerate library, although this may change in the future. For dataset loading, we use webdataset library. So the dataset should be in the webdataset format.

Models

At the momemnt we support following models:

The models are implemented under muse directory. All models implement the familiar transformers API. So you can use from_pretrained and save_pretrained methods to load and save the models. The model can be saved and loaded from the huggingface hub.

VQGAN example:

import torch from torchvision import transforms from PIL import Image from muse import MaskGitVQGAN

Load the pre-trained vq model from the hub

vq_model = MaskGitVQGAN.from_pretrained("openMUSE/maskgit-vqgan-imagenet-f16-256")

encode and decode images using

encode_transform = = transforms.Compose( [ transforms.Resize(256, interpolation=transforms.InterpolationMode.BILINEAR), transforms.CenterCrop(256), transforms.ToTensor(), ] ) image = Image.open("...") # pixel_values = encode_transform(image).unsqueeze(0) image_tokens = vq_model.encode(pixel_values) rec_image = vq_model.decode(image_tokens)

Convert to PIL images

rec_image = 2.0 * rec_image - 1.0 rec_image = torch.clamp(rec_image, -1.0, 1.0) rec_image = (rec_image + 1.0) / 2.0 rec_image *= 255.0 rec_image = rec_image.permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8) pil_images = [Image.fromarray(image) for image in rec_image]

MaskGitTransformer example for class-conditional generation:

import torch from muse import MaskGitTransformer, MaskGitVQGAN from muse.sampling import cosine_schedule

Load the pre-trained vq model from the hub

vq_model = MaskGitVQGAN.from_pretrained("openMUSE/maskgit-vqgan-imagenet-f16-256")

Initialize the MaskGitTransformer model

maskgit_model = MaskGitTransformer( vocab_size=2025, #(1024 + 1000 + 1 = 2025 -> Vq_tokens + Imagenet class ids + ) max_position_embeddings=257, # 256 + 1 for class token hidden_size=512, num_hidden_layers=8, num_attention_heads=8, intermediate_size=2048, codebook_size=1024, num_vq_tokens=256, num_classes=1000, )

prepare the input batch

images = torch.randn(4, 3, 256, 256) class_ids = torch.randint(0, 1000, (4,)) # random class ids

encode the images

image_tokens = vq_model.encode(images) batch_size, seq_len = image_tokens.shape

Sample a random timestep for each image

timesteps = torch.rand(batch_size, device=image_tokens.device)

Sample a random mask probability for each image using timestep and cosine schedule

mask_prob = cosine_schedule(timesteps) mask_prob = mask_prob.clip(min_masking_rate)

creat a random mask for each image

num_token_masked = (seq_len * mask_prob).round().clamp(min=1) batch_randperm = torch.rand(batch_size, seq_len, device=image_tokens.device).argsort(dim=-1) mask = batch_randperm < num_token_masked.unsqueeze(-1)

mask images and create input and labels

input_ids = torch.where(mask, mask_id, image_tokens) labels = torch.where(mask, image_tokens, -100)

shift the class ids by codebook size

class_ids = class_ids + vq_model.num_embeddings

prepend the class ids to the image tokens

input_ids = torch.cat([class_ids.unsqueeze(-1), input_ids], dim=-1)

prepend -100 to the labels as we don't want to predict the class ids

labels = torch.cat([-100 * torch.ones_like(class_ids).unsqueeze(-1), labels], dim=-1)

forward pass

logits, loss = maskgit_model(input_ids, labels=labels)

loss.backward()

to generate images

class_ids = torch.randint(0, 1000, (4,)) # random class ids generated_tokens = maskgit_model.generate(class_ids=class_ids) rec_images = vq_model.decode(generated_tokens)

Note:

Basic explanation of MaskGit Generation Process

  1. Maskgits is a transformer that outputs logits given a sequence of tokens of both vq and class-conditioned label token
  2. The way the denoising process is done is to mask out with mask token ids and gradually denoise
  3. In the original implementation, this is done by first using a softmax on the last dim and randomly sampling as a categorical distribution. This will give our predicted tokens for each maskid. Then we get the probabilities for those tokens to be chosen. Finally, we get the topk highest confidence probabilities when gumbel*temp is added to it. Gumbel distribution is like a shifted normal distribution towards 0 which is used to model extreme events. So in extreme scenarios, we will like to see a different token being chosen from the default one
  4. For the lucidrian implementation, it first removes the highest-scoring (lowest probability) tokens by masking them with a given masking ratio. Then, except for the highest 10% of the logits that we get, we set it to -infinity so when we do the gumbel distribution on it, they will be ignored. Then update the input ids and the scores where the scores are just 1-the probability given by the softmax of the logits at the predicted ids interestingly

Training

For class-conditional imagenet we are using accelerate for DDP training and webdataset for data loading. The training script is available in training/train_maskgit_imagenet.py.

We use OmegaConf for configuration management. See configs/template_config.yaml for the configuration template. Below we explain the configuration parameters.

wandb: entity: ???

experiment: name: ??? project: ??? output_dir: ??? max_train_examples: ??? save_every: 1000 eval_every: 500 generate_every: 1000 log_every: 50 log_grad_norm_every: 100 resume_from_checkpoint: latest

model: vq_model: pretrained: "openMUSE/maskgit-vqgan-imagenet-f16-256"

transformer:
    vocab_size: 2048 # (1024 + 1000 + 1 = 2025 -> Vq + Imagenet + <mask>, use 2048 for even division by 8)
    max_position_embeddings: 264 # (256 + 1 for class id, use 264 for even division by 8)
    hidden_size: 768
    num_hidden_layers: 12
    num_attention_heads: 12
    intermediate_size: 3072
    codebook_size: 1024
    num_vq_tokens: 256
    num_classes: 1000
    initializer_range: 0.02
    layer_norm_eps: 1e-6
    use_bias: False
    use_normformer: True
    use_encoder_layernorm: True
    hidden_dropout: 0.0
    attention_dropout: 0.0

gradient_checkpointing: True
enable_xformers_memory_efficient_attention: False

dataset: params: train_shards_path_or_url: ??? eval_shards_path_or_url: ??? batch_size: ${training.batch_size} shuffle_buffer_size: ??? num_workers: ??? resolution: 256 pin_memory: True persistent_workers: True preprocessing: resolution: 256 center_crop: True random_flip: False

optimizer: name: adamw # Can be adamw or lion or fused_adamw. Install apex for fused_adamw params: # default adamw params learning_rate: ??? scale_lr: False # scale learning rate by total batch size beta1: 0.9 beta2: 0.999 weight_decay: 0.01 epsilon: 1e-8

lr_scheduler: scheduler: "constant_with_warmup" params: learning_rate: ${optimizer.params.learning_rate} warmup_steps: 500

training: gradient_accumulation_steps: 1 batch_size: 128 mixed_precision: "no" enable_tf32: True use_ema: False seed: 42 max_train_steps: ??? overfit_one_batch: False min_masking_rate: 0.0 label_smoothing: 0.0 max_grad_norm: null

Arguments with ??? are required.

wandb:

experiment:

model:

dataset:

optimizer:

lr_scheduler:

training:

Notes about training and dataset.:

We randomly resample the shards (with replacement) and sample examples in buffer for training every time we resume/start the training run. This means our data loading is not determinitsic. We also don't do epoch based training but just using this for book keeping and being able to reuse the same training loop with other datasets/loaders.

Running experiments:

So far we are running experiments on single node. To launch a training run on a single node, run the following steps:

  1. Prepare the dataset in webdataset format. You can use the scripts/convert_imagenet_to_wds.py script to convert the imagenet dataset to webdataset format.
  2. First configure your training env using accelerate config.
  3. Create a config.yaml file for your experiment.
  4. Launch the training run using accelerate launch.

accelerate launch python -u training/train_maskgit_imagenet.py config=path/to/yaml/config

With OmegaConf, commandline overrides are done in dot-notation format. E.g. if you want to override the dataset path, you would use the command python -u train.py config=path/to/config dataset.params.path=path/to/dataset.

The same command can be used to launch the training locally.

Steps

Setup the codebase and train a class-conditional model no imagenet.

Conduct text2image experiments on CC12M.

Train improved VQGANs models.

Misc tasks

Repo structure (WIP)

├── README.md
├── configs                        -> All training config files.
│   └── template_config.yaml
├── muse
│   ├── __init__.py
│   ├── data.py                    -> All data related utils. Can create a data folder if needed.
│   ├── logging.py                 -> Misc logging utils.
|   ├── lr_schedulers.py           -> All lr scheduler related utils.
│   ├── modeling_maskgit_vqgan.py  -> VQGAN model from maskgit repo.
│   ├── modeling_taming_vqgan.py   -> VQGAN model from taming repo.
│   └── modeling_transformer.py    -> The main transformer model.
│   ├── modeling_utils.py          -> All model related utils, like save_pretrained, from_pretrained from hub etc
│   ├── sampling.py                -> Sampling/Generation utils.
│   ├── training_utils.py          -> Common training utils.
├── pyproject.toml
├── setup.cfg
├── setup.py
├── test.py
└── training                       -> All training scripts.
    ├── __init__.py
    ├── data.py                    -> All data related utils. Can create a data folder if needed.
    ├── optimizer.py               -> All optimizer related utils and any new optimizer not available in PT.
    ├── train_maskgit_imagenet.py
    ├── train_muse.py
    └── train_vqgan.py

Acknowledgments

This project is hevily based on the following open-source repos. Thanks to all the authors for their amazing work.

And obivioulsy to PyTorch team for this amazing framework ❤️