Package: taming-transformers-hugf
Taming Transformers for High-Resolution Image Synthesisnote: forked by tongxin.wen for easier use in hugging-face, for example, added the from_pretrained mixin for modelsCVPR 2021 (Oral)Taming Transformers for High-Resolution Image SynthesisPatrick Esser*,Robin Rombach*,Björn Ommer* equal contributiontl;drWe combine the efficiancy of convolutional approaches with the expressivity of transformers by introducing a convolutional VQGAN, which learns a codebook of context-rich visual parts, whose composition is modeled with an autoregressive transformer.arXiv|BibTeX|Project PageNews2022More pretrained VQGANs (e.g. a f8-model with only 256 codebook entries) are available in our new work onLatent Diffusion Models.Added scene synthesis models as proposed in the paperHigh-Resolution Complex Scene Synthesis with Transformers, seethis section.2021Thanks torom1504it is now easy totrain a VQGAN on your own datasets.Included a bugfix for the quantizer. For backward compatibility it is
disabled by default (which corresponds to always training withbeta=1.0).
Uselegacy=Falsein the quantizer config to enable it.
Thanksrichcmwangandwcshin-git!Our paper received an update: Seehttps://arxiv.org/abs/2012.09841v3and the corresponding changelog.Added a pretrained,1.4B transformer modeltrained for class-conditional ImageNet synthesis, which obtains state-of-the-art FID scores among autoregressive approaches and outperforms BigGAN.Added pretrained, unconditional models onFFHQandCelebA-HQ.Added accelerated sampling via caching of keys/values in the self-attention operation, used inscripts/sample_fast.py.Added a checkpoint of aVQGANtrained with f8 compression and Gumbel-Quantization.
See also our updatedreconstruction notebook.We added acolab notebookwhich compares two VQGANs and OpenAI'sDALL-E. See alsothis section.We now include an overview of pretrained models inTab.1. We added models forCOCOandADE20k.The streamlit demo now supports image completions.We now include a couple of examples from the D-RIN dataset so you can run theD-RIN demowithout preparing the dataset first.You can now jump right into sampling with ourColab quickstart notebook.RequirementsA suitablecondaenvironment namedtamingcan be created
and activated with:conda env create -f environment.yaml
conda activate tamingOverview of pretrained modelsThe following table provides an overview of all models that are currently available.
FID scores were evaluated usingtorch-fidelity.
For reference, we also include a link to the recently released autoencoder of theDALL-Emodel.
See the correspondingcolab
notebookfor a comparison and discussion of reconstruction capabilities.DatasetFID vs trainFID vs valLinkSamples (256x256)CommentsFFHQ (f=16)9.6--ffhq_transformerffhq_samplesCelebA-HQ (f=16)10.2--celebahq_transformercelebahq_samplesADE20K (f=16)--35.5ade20k_transformerade20k_samples.zip[2k]evaluated on val split (2k images)COCO-Stuff (f=16)--20.4coco_transformercoco_samples.zip[5k]evaluated on val split (5k images)ImageNet (cIN) (f=16)15.98/15.78/6.59/5.88/5.20--cin_transformercin_samplesdifferent decoding hyperparametersFacesHQ (f=16)----faceshq_transformerS-FLCKR (f=16)----sflckrD-RIN (f=16)----drin_transformerVQGAN ImageNet (f=16), 102410.547.94vqgan_imagenet_f16_1024reconstructionsReconstruction-FIDs.VQGAN ImageNet (f=16), 163847.414.98vqgan_imagenet_f16_16384reconstructionsReconstruction-FIDs.VQGAN OpenImages (f=8), 256--1.49https://ommer-lab.com/files/latent-diffusion/vq-f8-n256.zip---Reconstruction-FIDs. Available vialatent diffusion.VQGAN OpenImages (f=8), 16384--1.14https://ommer-lab.com/files/latent-diffusion/vq-f8.zip---Reconstruction-FIDs. Available vialatent diffusionVQGAN OpenImages (f=8), 8192, GumbelQuantization3.241.49vqgan_gumbel_f8---Reconstruction-FIDs.DALL-E dVAE (f=8), 8192, GumbelQuantization33.8832.01https://github.com/openai/DALL-EreconstructionsReconstruction-FIDs.Running pretrained modelsThe commands below will start a streamlit demo which supports sampling at
different resolutions and image completions. To run a non-interactive version
of the sampling process, replacestreamlit run scripts/sample_conditional.py --bypython scripts/make_samples.py --outdir <path_to_write_samples_to>and
keep the remaining command line arguments.To sample from unconditional or class-conditional models,
runpython scripts/sample_fast.py -r <path/to/config_and_checkpoint>.
We describe below how to use this script to sample from the ImageNet, FFHQ, and CelebA-HQ models,
respectively.S-FLCKRYou can alsorun this model in a Colab
notebook,
which includes all necessary steps to start sampling.Download the2020-11-09T13-31-51_sflckrfolder and place it intologs. Then, runstreamlit run scripts/sample_conditional.py -- -r logs/2020-11-09T13-31-51_sflckr/ImageNetDownload the2021-04-03T19-39-50_cin_transformerfolder and place it into logs.  Sampling from the class-conditional ImageNet
model does not require any data preparation. To produce 50 samples for each of
the 1000 classes of ImageNet, with k=600 for top-k sampling, p=0.92 for nucleus
sampling and temperature t=1.0, runpython scripts/sample_fast.py -r logs/2021-04-03T19-39-50_cin_transformer/ -n 50 -k 600 -t 1.0 -p 0.92 --batch_size 25To restrict the model to certain classes, provide them via the--classesargument, separated by
commas. For example, to sample 50ostriches,border colliesandwhiskey jugs, runpython scripts/sample_fast.py -r logs/2021-04-03T19-39-50_cin_transformer/ -n 50 -k 600 -t 1.0 -p 0.92 --batch_size 25 --classes 9,232,901We recommended to experiment with the autoregressive decoding parameters (top-k, top-p and temperature) for best results.FFHQ/CelebA-HQDownload the2021-04-23T18-19-01_ffhq_transformerand2021-04-23T18-11-19_celebahq_transformerfolders and place them into logs.
Again, sampling from these unconditional models does not require any data preparation.
To produce 50000 samples, with k=250 for top-k sampling,
p=1.0 for nucleus sampling and temperature t=1.0, runpython scripts/sample_fast.py -r logs/2021-04-23T18-19-01_ffhq_transformer/for FFHQ andpython scripts/sample_fast.py -r logs/2021-04-23T18-11-19_celebahq_transformer/to sample from the CelebA-HQ model.
For both models it can be advantageous to vary the top-k/top-p parameters for sampling.FacesHQDownload2020-11-13T21-41-45_faceshq_transformerand
place it intologs. Follow the data preparation steps forCelebA-HQandFFHQ. Runstreamlit run scripts/sample_conditional.py -- -r logs/2020-11-13T21-41-45_faceshq_transformer/D-RINDownload2020-11-20T12-54-32_drin_transformerand
place it intologs. To run the demo on a couple of example depth maps
included in the repository, runstreamlit run scripts/sample_conditional.py -- -r logs/2020-11-20T12-54-32_drin_transformer/ --ignore_base_data data="{target: main.DataModuleFromConfig, params: {batch_size: 1, validation: {target: taming.data.imagenet.DRINExamples}}}"To run the demo on the complete validation set, first follow the data preparation steps forImageNetand then runstreamlit run scripts/sample_conditional.py -- -r logs/2020-11-20T12-54-32_drin_transformer/COCODownload2021-01-20T16-04-20_coco_transformerand
place it intologs. To run the demo on a couple of example segmentation maps
included in the repository, runstreamlit run scripts/sample_conditional.py -- -r logs/2021-01-20T16-04-20_coco_transformer/ --ignore_base_data data="{target: main.DataModuleFromConfig, params: {batch_size: 1, validation: {target: taming.data.coco.Examples}}}"ADE20kDownload2020-11-20T21-45-44_ade20k_transformerand
place it intologs. To run the demo on a couple of example segmentation maps
included in the repository, runstreamlit run scripts/sample_conditional.py -- -r logs/2020-11-20T21-45-44_ade20k_transformer/ --ignore_base_data data="{target: main.DataModuleFromConfig, params: {batch_size: 1, validation: {target: taming.data.ade20k.Examples}}}"Scene Image SynthesisScene image generation based on bounding box conditionals as done in our CVPR2021 AI4CC workshop paperHigh-Resolution Complex Scene Synthesis with Transformers(see talk onworkshop page). Supporting the datasets COCO and Open Images.TrainingDownload first-stage modelsCOCO-8k-VQGANfor COCO orCOCO/Open-Images-8k-VQGANfor Open Images.
Changeckpt_pathindata/coco_scene_images_transformer.yamlanddata/open_images_scene_images_transformer.yamlto point to the downloaded first-stage models.
Download the full COCO/OI datasets and adaptdata_pathin the same files, unless working with the 100 files provided for training and validation suits your needs already.Code can be run withpython main.py --base configs/coco_scene_images_transformer.yaml -t True --gpus 0,orpython main.py --base configs/open_images_scene_images_transformer.yaml -t True --gpus 0,SamplingTrain a model as described above or download a pre-trained model:Open Images 1 billion parameter modelavailable that trained 100 epochs. On 256x256 pixels, FID 41.48±0.21, SceneFID 14.60±0.15, Inception Score 18.47±0.27. The model was trained with 2d crops of images and is thus well-prepared for the task of generating high-resolution images, e.g. 512x512.Open Images distilled version of the above model with 125 million parametersallows for sampling on smaller GPUs (4 GB is enough for sampling 256x256 px images). Model was trained for 60 epochs with 10% soft loss, 90% hard loss. On 256x256 pixels, FID 43.07±0.40, SceneFID 15.93±0.19, Inception Score 17.23±0.11.COCO 30 epochsCOCO 60 epochs(find model statistics for both COCO versions inassets/coco_scene_images_training.svg)When downloading a pre-trained model, remember to changeckpt_pathinconfigs/*project.yamlto point to your downloaded first-stage model (see ->Training).Scene image generation can be run withpython scripts/make_scene_samples.py --outdir=/some/outdir -r /path/to/pretrained/model --resolution=512,512Training on custom dataTraining on your own dataset can be beneficial to get better tokens and hence better images for your domain.
Those are the steps to follow to make this work:install the repo withconda env create -f environment.yaml,conda activate tamingandpip install -e .put your .jpg files in a folderyour_foldercreate 2 text files axx_train.txtandxx_test.txtthat point to the files in your training and test set respectively (for examplefind $(pwd)/your_folder -name "*.jpg" > train.txt)adaptconfigs/custom_vqgan.yamlto point to these 2 filesrunpython main.py --base configs/custom_vqgan.yaml -t True --gpus 0,1to
train on two GPUs. Use--gpus 0,(with a trailing comma) to train on a single GPU.Data PreparationImageNetThe code will try to download (throughAcademic
Torrents) and prepare ImageNet the first time it
is used. However, since ImageNet is quite large, this requires a lot of disk
space and time. If you already have ImageNet on your disk, you can speed things
up by putting the data into${XDG_CACHE}/autoencoders/data/ILSVRC2012_{split}/data/(which defaults to~/.cache/autoencoders/data/ILSVRC2012_{split}/data/), where{split}is one
oftrain/validation. It should have the following structure:${XDG_CACHE}/autoencoders/data/ILSVRC2012_{split}/data/
├── n01440764
│   ├── n01440764_10026.JPEG
│   ├── n01440764_10027.JPEG
│   ├── ...
├── n01443537
│   ├── n01443537_10007.JPEG
│   ├── n01443537_10014.JPEG
│   ├── ...
├── ...If you haven't extracted the data, you can also placeILSVRC2012_img_train.tar/ILSVRC2012_img_val.tar(or symlinks to them) into${XDG_CACHE}/autoencoders/data/ILSVRC2012_train//${XDG_CACHE}/autoencoders/data/ILSVRC2012_validation/, which will then be
extracted into above structure without downloading it again.  Note that this
will only happen if neither a folder${XDG_CACHE}/autoencoders/data/ILSVRC2012_{split}/data/nor a file${XDG_CACHE}/autoencoders/data/ILSVRC2012_{split}/.readyexist. Remove them
if you want to force running the dataset preparation again.You will then need to prepare the depth data usingMiDaS. Create a symlinkdata/imagenet_depthpointing to a folder with two subfolderstrainandval, each mirroring the structure of the corresponding ImageNet folder
described above and containing apngfile for each of ImageNet'sJPEGfiles. Thepngencodesfloat32depth values obtained from MiDaS as RGBA
images. We provide the scriptscripts/extract_depth.pyto generate this data.Please notethat this script usesMiDaS via PyTorch
Hub. When we prepared the data,
the hub provided theMiDaS
v2.0version, but now it
provides a v2.1 version. We haven't tested our models with depth maps obtained
via v2.1 and if you want to make sure that things work as expected, you must
adjust the script to make sure it explicitly usesv2.0!CelebA-HQCreate a symlinkdata/celebahqpointing to a folder containing the.npyfiles of CelebA-HQ (instructions to obtain them can be found in thePGGAN
repository).FFHQCreate a symlinkdata/ffhqpointing to theimages1024x1024folder obtained
from theFFHQ repository.S-FLCKRUnfortunately, we are not allowed to distribute the images we collected for the
S-FLCKR dataset and can therefore only give a description how it was produced.
There are many resources oncollecting images from the
webto get started.
We collected sufficiently large images fromflickr(seedata/flickr_tags.txtfor a full list of tags used to find images)
and varioussubreddits(seedata/subreddits.txtfor all subreddits that were used).
Overall, we collected 107625 images, and split them randomly into 96861
training images and 10764 validation images. We then obtained segmentation
masks for each image usingDeepLab v2trained onCOCO-Stuff. We used aPyTorch
reimplementationand include an
example script for this process inscripts/extract_segmentation.py.COCOCreate a symlinkdata/cococontaining the images from the 2017 split intrain2017andval2017, and their annotations inannotations. Files can be
obtained from theCOCO webpage. In addition, we use
theStuff+thing PNG-style annotations on COCO 2017
trainvalannotations fromCOCO-Stuff, which
should be placed underdata/cocostuffthings.ADE20kCreate a symlinkdata/ade20k_rootcontaining the contents ofADEChallengeData2016.zipfrom theMIT Scene Parsing Benchmark.Training modelsFacesHQTrain a VQGAN withpython main.py --base configs/faceshq_vqgan.yaml -t True --gpus 0,Then, adjust the checkpoint path of the config keymodel.params.first_stage_config.params.ckpt_pathinconfigs/faceshq_transformer.yaml(or download2020-11-09T13-33-36_faceshq_vqganand place intologs, which
corresponds to the preconfigured checkpoint path), then runpython main.py --base configs/faceshq_transformer.yaml -t True --gpus 0,D-RINTrain a VQGAN on ImageNet withpython main.py --base configs/imagenet_vqgan.yaml -t True --gpus 0,or download a pretrained one from2020-09-23T17-56-33_imagenet_vqganand place underlogs. If you trained your own, adjust the path in the config
keymodel.params.first_stage_config.params.ckpt_pathofconfigs/drin_transformer.yaml.Train a VQGAN on Depth Maps of ImageNet withpython main.py --base configs/imagenetdepth_vqgan.yaml -t True --gpus 0,or download a pretrained one from2020-11-03T15-34-24_imagenetdepth_vqganand place underlogs. If you trained your own, adjust the path in the config
keymodel.params.cond_stage_config.params.ckpt_pathofconfigs/drin_transformer.yaml.To train the transformer, runpython main.py --base configs/drin_transformer.yaml -t True --gpus 0,More ResourcesComparing Different First Stage ModelsThe reconstruction and compression capabilities of different fist stage models can be analyzed in thiscolab notebook.
In particular, the notebook compares two VQGANs with a downsampling factor of f=16 for each and codebook dimensionality of 1024 and 16384,
a VQGAN with f=8 and 8192 codebook entries and the discrete autoencoder of OpenAI'sDALL-E(which has f=8 and 8192
codebook entries).OtherAvideo summarybyTwo Minute Papers.Avideo summarybyGradient Dude.Aweights and biases report summarizing the paperbyayulockin.Avideo summarybyWhat's AI.Take a look atak9250's notebookif you want to run the streamlit demos on Colab.Text-to-Image Optimization via CLIPVQGAN has been successfully used as an image generator guided by theCLIPmodel, both for pure image generation
from scratch and image-to-image translation. We recommend the following notebooks/videos/resources:AdvadnounsPatreon and corresponding LatentVision notebooks:https://www.patreon.com/patronizemeThenotebookofRivers Have Wings.Avideoexplanation byDot CSV(in Spanish, but English subtitles are available)Text prompt:'A bird drawn by a child'Shout-outsThanks to everyone who makes their code and models available. In particular,The architecture of our VQGAN is inspired byDenoising Diffusion Probabilistic ModelsThe very hackable transformer implementationminGPTThe good ol'PatchGANandLearned Perceptual Similarity (LPIPS)BibTeX@misc{esser2020taming,
      title={Taming Transformers for High-Resolution Image Synthesis}, 
      author={Patrick Esser and Robin Rombach and Björn Ommer},
      year={2020},
      eprint={2012.09841},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}