test rm taming

taming-transformers/.gitignore

@@ -1,2 +0,0 @@
-assets/
-data/

taming-transformers/License.txt

@@ -1,19 +0,0 @@
-Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer
-
-Permission is hereby granted, free of charge, to any person obtaining a copy
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
-
-The above copyright notice and this permission notice shall be included in all
-copies or substantial portions of the Software.
-
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
-EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
-MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
-IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
-DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
-OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
-OR OTHER DEALINGS IN THE SOFTWARE./

taming-transformers/README.md

@@ -1,410 +0,0 @@
-# Taming Transformers for High-Resolution Image Synthesis
-##### CVPR 2021 (Oral)
-![teaser](assets/mountain.jpeg)
-
-[**Taming Transformers for High-Resolution Image Synthesis**](https://compvis.github.io/taming-transformers/)<br/>
-[Patrick Esser](https://github.com/pesser)\*,
-[Robin Rombach](https://github.com/rromb)\*,
-[Björn Ommer](https://hci.iwr.uni-heidelberg.de/Staff/bommer)<br/>
-\* equal contribution
-
-**tl;dr** We 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.
-
-![teaser](assets/teaser.png)
-[arXiv](https://arxiv.org/abs/2012.09841) | [BibTeX](#bibtex) | [Project Page](https://compvis.github.io/taming-transformers/)
-
-
-### News
-#### 2022
-- More pretrained VQGANs (e.g. a f8-model with only 256 codebook entries) are available in our new work on [Latent Diffusion Models](https://github.com/CompVis/latent-diffusion).
-- Added scene synthesis models as proposed in the paper [High-Resolution Complex Scene Synthesis with Transformers](https://arxiv.org/abs/2105.06458), see [this section](#scene-image-synthesis).
-#### 2021
-- Thanks to [rom1504](https://github.com/rom1504) it is now easy to [train a VQGAN on your own datasets](#training-on-custom-data).
-- Included a bugfix for the quantizer. For backward compatibility it is
-  disabled by default (which corresponds to always training with `beta=1.0`).
-  Use `legacy=False` in the quantizer config to enable it.
-  Thanks [richcmwang](https://github.com/richcmwang) and [wcshin-git](https://github.com/wcshin-git)!
-- Our paper received an update: See https://arxiv.org/abs/2012.09841v3 and the corresponding changelog.
-- Added a pretrained, [1.4B transformer model](https://k00.fr/s511rwcv) trained for class-conditional ImageNet synthesis, which obtains state-of-the-art FID scores among autoregressive approaches and outperforms BigGAN.
-- Added pretrained, unconditional models on [FFHQ](https://k00.fr/yndvfu95) and [CelebA-HQ](https://k00.fr/2xkmielf).
-- Added accelerated sampling via caching of keys/values in the self-attention operation, used in `scripts/sample_fast.py`.
-- Added a checkpoint of a [VQGAN](https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/) trained with f8 compression and Gumbel-Quantization. 
-  See also our updated [reconstruction notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/reconstruction_usage.ipynb). 
-- We added a [colab notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/reconstruction_usage.ipynb) which compares two VQGANs and OpenAI's [DALL-E](https://github.com/openai/DALL-E). See also [this section](#more-resources).
-- We now include an overview of pretrained models in [Tab.1](#overview-of-pretrained-models). We added models for [COCO](#coco) and [ADE20k](#ade20k).
-- The streamlit demo now supports image completions.
-- We now include a couple of examples from the D-RIN dataset so you can run the
-  [D-RIN demo](#d-rin) without preparing the dataset first.
-- You can now jump right into sampling with our [Colab quickstart notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/taming-transformers.ipynb).
-
-## Requirements
-A suitable [conda](https://conda.io/) environment named `taming` can be created
-and activated with:
-
-```
-conda env create -f environment.yaml
-conda activate taming
-```
-## Overview of pretrained models
-The following table provides an overview of all models that are currently available. 
-FID scores were evaluated using [torch-fidelity](https://github.com/toshas/torch-fidelity).
-For reference, we also include a link to the recently released autoencoder of the [DALL-E](https://github.com/openai/DALL-E) model. 
-See the corresponding [colab
-notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/reconstruction_usage.ipynb)
-for a comparison and discussion of reconstruction capabilities.
-
-| Dataset  | FID vs train | FID vs val | Link |  Samples (256x256) | Comments
-| ------------- | ------------- | ------------- |-------------  | -------------  |-------------  |
-| FFHQ (f=16) | 9.6 | -- | [ffhq_transformer](https://k00.fr/yndvfu95) |  [ffhq_samples](https://k00.fr/j626x093) |
-| CelebA-HQ (f=16) | 10.2 | -- | [celebahq_transformer](https://k00.fr/2xkmielf) | [celebahq_samples](https://k00.fr/j626x093) |
-| ADE20K (f=16) | -- | 35.5 | [ade20k_transformer](https://k00.fr/ot46cksa) | [ade20k_samples.zip](https://heibox.uni-heidelberg.de/f/70bb78cbaf844501b8fb/) [2k] | evaluated on val split (2k images)
-| COCO-Stuff (f=16) | -- | 20.4  | [coco_transformer](https://k00.fr/2zz6i2ce) | [coco_samples.zip](https://heibox.uni-heidelberg.de/f/a395a9be612f4a7a8054/) [5k] | evaluated on val split (5k images)
-| ImageNet (cIN) (f=16) | 15.98/15.78/6.59/5.88/5.20 | -- | [cin_transformer](https://k00.fr/s511rwcv) | [cin_samples](https://k00.fr/j626x093) | different decoding hyperparameters |  
-| |  | | || |
-| FacesHQ (f=16) | -- |  -- | [faceshq_transformer](https://k00.fr/qqfl2do8)
-| S-FLCKR (f=16) | -- | -- | [sflckr](https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/) 
-| D-RIN (f=16) | -- | -- | [drin_transformer](https://k00.fr/39jcugc5)
-| | |  | | || |
-| VQGAN ImageNet (f=16), 1024 |  10.54 | 7.94 | [vqgan_imagenet_f16_1024](https://heibox.uni-heidelberg.de/d/8088892a516d4e3baf92/) | [reconstructions](https://k00.fr/j626x093) | Reconstruction-FIDs.
-| VQGAN ImageNet (f=16), 16384 | 7.41 | 4.98 |[vqgan_imagenet_f16_16384](https://heibox.uni-heidelberg.de/d/a7530b09fed84f80a887/)  |  [reconstructions](https://k00.fr/j626x093) | Reconstruction-FIDs.
-| VQGAN OpenImages (f=8), 256 | -- | 1.49 |https://ommer-lab.com/files/latent-diffusion/vq-f8-n256.zip |  ---  | Reconstruction-FIDs. Available via [latent diffusion](https://github.com/CompVis/latent-diffusion).
-| VQGAN OpenImages (f=8), 16384 | -- | 1.14 |https://ommer-lab.com/files/latent-diffusion/vq-f8.zip  |  ---  | Reconstruction-FIDs. Available via [latent diffusion](https://github.com/CompVis/latent-diffusion)
-| VQGAN OpenImages (f=8), 8192, GumbelQuantization | 3.24 | 1.49 |[vqgan_gumbel_f8](https://heibox.uni-heidelberg.de/d/2e5662443a6b4307b470/)  |  ---  | Reconstruction-FIDs.
-| | |  | | || |
-| DALL-E dVAE (f=8), 8192, GumbelQuantization | 33.88 | 32.01 | https://github.com/openai/DALL-E | [reconstructions](https://k00.fr/j626x093) | Reconstruction-FIDs.
-
-
-## Running pretrained models
-
-The 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, replace `streamlit run scripts/sample_conditional.py --`
-by `python 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, 
-run `python 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-FLCKR
-![teaser](assets/sunset_and_ocean.jpg)
-
-You can also [run this model in a Colab
-notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/taming-transformers.ipynb),
-which includes all necessary steps to start sampling.
-
-Download the
-[2020-11-09T13-31-51_sflckr](https://heibox.uni-heidelberg.de/d/73487ab6e5314cb5adba/)
-folder and place it into `logs`. Then, run
-```
-streamlit run scripts/sample_conditional.py -- -r logs/2020-11-09T13-31-51_sflckr/
-```
-
-### ImageNet
-![teaser](assets/imagenet.png)
-
-Download the [2021-04-03T19-39-50_cin_transformer](https://k00.fr/s511rwcv)
-folder 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, run
-
-```
-python 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   
-```
-
-To restrict the model to certain classes, provide them via the `--classes` argument, separated by 
-commas. For example, to sample 50 *ostriches*, *border collies* and *whiskey jugs*, run
-
-```
-python 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,901   
-```
-We recommended to experiment with the autoregressive decoding parameters (top-k, top-p and temperature) for best results.  
-
-### FFHQ/CelebA-HQ
-
-Download the [2021-04-23T18-19-01_ffhq_transformer](https://k00.fr/yndvfu95) and 
-[2021-04-23T18-11-19_celebahq_transformer](https://k00.fr/2xkmielf) 
-folders 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, run
-
-```
-python scripts/sample_fast.py -r logs/2021-04-23T18-19-01_ffhq_transformer/   
-```
-for FFHQ and  
-
-```
-python 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.
-
-### FacesHQ
-![teaser](assets/faceshq.jpg)
-
-Download [2020-11-13T21-41-45_faceshq_transformer](https://k00.fr/qqfl2do8) and
-place it into `logs`. Follow the data preparation steps for
-[CelebA-HQ](#celeba-hq) and [FFHQ](#ffhq). Run
-```
-streamlit run scripts/sample_conditional.py -- -r logs/2020-11-13T21-41-45_faceshq_transformer/
-```
-
-### D-RIN
-![teaser](assets/drin.jpg)
-
-Download [2020-11-20T12-54-32_drin_transformer](https://k00.fr/39jcugc5) and
-place it into `logs`. To run the demo on a couple of example depth maps
-included in the repository, run
-
-```
-streamlit 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 for
-[ImageNet](#imagenet) and then run
-```
-streamlit run scripts/sample_conditional.py -- -r logs/2020-11-20T12-54-32_drin_transformer/
-```
-
-### COCO
-Download [2021-01-20T16-04-20_coco_transformer](https://k00.fr/2zz6i2ce) and
-place it into `logs`. To run the demo on a couple of example segmentation maps
-included in the repository, run
-
-```
-streamlit 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}}}"
-```
-
-### ADE20k
-Download [2020-11-20T21-45-44_ade20k_transformer](https://k00.fr/ot46cksa) and
-place it into `logs`. To run the demo on a couple of example segmentation maps
-included in the repository, run
-
-```
-streamlit 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 Synthesis
-![teaser](assets/scene_images_samples.svg)
-Scene image generation based on bounding box conditionals as done in our CVPR2021 AI4CC workshop paper [High-Resolution Complex Scene Synthesis with Transformers](https://arxiv.org/abs/2105.06458) (see talk on [workshop page](https://visual.cs.brown.edu/workshops/aicc2021/#awards)). Supporting the datasets COCO and Open Images.
-
-### Training
-Download first-stage models [COCO-8k-VQGAN](https://heibox.uni-heidelberg.de/f/78dea9589974474c97c1/) for COCO or [COCO/Open-Images-8k-VQGAN](https://heibox.uni-heidelberg.de/f/461d9a9f4fcf48ab84f4/) for Open Images.
-Change `ckpt_path` in `data/coco_scene_images_transformer.yaml` and `data/open_images_scene_images_transformer.yaml` to point to the downloaded first-stage models.
-Download the full COCO/OI datasets and adapt `data_path` in the same files, unless working with the 100 files provided for training and validation suits your needs already.
-
-Code can be run with
-`python main.py --base configs/coco_scene_images_transformer.yaml -t True --gpus 0,`
-or
-`python main.py --base configs/open_images_scene_images_transformer.yaml -t True --gpus 0,`
-
-### Sampling 
-Train a model as described above or download a pre-trained model:
- - [Open Images 1 billion parameter model](https://drive.google.com/file/d/1FEK-Z7hyWJBvFWQF50pzSK9y1W_CJEig/view?usp=sharing) available 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 parameters](https://drive.google.com/file/d/1xf89g0mc78J3d8Bx5YhbK4tNRNlOoYaO) allows 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 epochs](https://heibox.uni-heidelberg.de/f/0d0b2594e9074c7e9a33/)
- - [COCO 60 epochs](https://drive.google.com/file/d/1bInd49g2YulTJBjU32Awyt5qnzxxG5U9/) (find model statistics for both COCO versions in `assets/coco_scene_images_training.svg`)
-
-When downloading a pre-trained model, remember to change `ckpt_path` in `configs/*project.yaml` to point to your downloaded first-stage model (see ->Training).
-
-Scene image generation can be run with
-`python scripts/make_scene_samples.py --outdir=/some/outdir -r /path/to/pretrained/model --resolution=512,512`
-
-
-## Training on custom data
-
-Training 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:
-1. install the repo with `conda env create -f environment.yaml`, `conda activate taming` and `pip install -e .`
-1. put your .jpg files in a folder `your_folder`
-2. create 2 text files a `xx_train.txt` and `xx_test.txt` that point to the files in your training and test set respectively (for example `find $(pwd)/your_folder -name "*.jpg" > train.txt`)
-3. adapt `configs/custom_vqgan.yaml` to point to these 2 files
-4. run `python main.py --base configs/custom_vqgan.yaml -t True --gpus 0,1` to
-   train on two GPUs. Use `--gpus 0,` (with a trailing comma) to train on a single GPU.
-
-## Data Preparation
-
-### ImageNet
-The code will try to download (through [Academic
-Torrents](http://academictorrents.com/)) 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
-of `train`/`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 place
-`ILSVRC2012_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}/.ready` exist. Remove them
-if you want to force running the dataset preparation again.
-
-You will then need to prepare the depth data using
-[MiDaS](https://github.com/intel-isl/MiDaS). Create a symlink
-`data/imagenet_depth` pointing to a folder with two subfolders `train` and
-`val`, each mirroring the structure of the corresponding ImageNet folder
-described above and containing a `png` file for each of ImageNet's `JPEG`
-files. The `png` encodes `float32` depth values obtained from MiDaS as RGBA
-images. We provide the script `scripts/extract_depth.py` to generate this data.
-**Please note** that this script uses [MiDaS via PyTorch
-Hub](https://pytorch.org/hub/intelisl_midas_v2/). When we prepared the data,
-the hub provided the [MiDaS
-v2.0](https://github.com/intel-isl/MiDaS/releases/tag/v2) version, 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 uses
-[v2.0](https://github.com/intel-isl/MiDaS/releases/tag/v2)!
-
-### CelebA-HQ
-Create a symlink `data/celebahq` pointing to a folder containing the `.npy`
-files of CelebA-HQ (instructions to obtain them can be found in the [PGGAN
-repository](https://github.com/tkarras/progressive_growing_of_gans)).
-
-### FFHQ
-Create a symlink `data/ffhq` pointing to the `images1024x1024` folder obtained
-from the [FFHQ repository](https://github.com/NVlabs/ffhq-dataset).
-
-### S-FLCKR
-Unfortunately, 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 on [collecting images from the
-web](https://github.com/adrianmrit/flickrdatasets) to get started.
-We collected sufficiently large images from [flickr](https://www.flickr.com)
-(see `data/flickr_tags.txt` for a full list of tags used to find images)
-and various [subreddits](https://www.reddit.com/r/sfwpornnetwork/wiki/network)
-(see `data/subreddits.txt` for 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 using [DeepLab v2](https://arxiv.org/abs/1606.00915)
-trained on [COCO-Stuff](https://arxiv.org/abs/1612.03716). We used a [PyTorch
-reimplementation](https://github.com/kazuto1011/deeplab-pytorch) and include an
-example script for this process in `scripts/extract_segmentation.py`.
-
-### COCO
-Create a symlink `data/coco` containing the images from the 2017 split in
-`train2017` and `val2017`, and their annotations in `annotations`. Files can be
-obtained from the [COCO webpage](https://cocodataset.org/). In addition, we use
-the [Stuff+thing PNG-style annotations on COCO 2017
-trainval](http://calvin.inf.ed.ac.uk/wp-content/uploads/data/cocostuffdataset/stuffthingmaps_trainval2017.zip)
-annotations from [COCO-Stuff](https://github.com/nightrome/cocostuff), which
-should be placed under `data/cocostuffthings`.
-
-### ADE20k
-Create a symlink `data/ade20k_root` containing the contents of
-[ADEChallengeData2016.zip](http://data.csail.mit.edu/places/ADEchallenge/ADEChallengeData2016.zip)
-from the [MIT Scene Parsing Benchmark](http://sceneparsing.csail.mit.edu/).
-
-## Training models
-
-### FacesHQ
-
-Train a VQGAN with
-```
-python main.py --base configs/faceshq_vqgan.yaml -t True --gpus 0,
-```
-
-Then, adjust the checkpoint path of the config key
-`model.params.first_stage_config.params.ckpt_path` in
-`configs/faceshq_transformer.yaml` (or download
-[2020-11-09T13-33-36_faceshq_vqgan](https://k00.fr/uxy5usa9) and place into `logs`, which
-corresponds to the preconfigured checkpoint path), then run
-```
-python main.py --base configs/faceshq_transformer.yaml -t True --gpus 0,
-```
-
-### D-RIN
-
-Train a VQGAN on ImageNet with
-```
-python main.py --base configs/imagenet_vqgan.yaml -t True --gpus 0,
-```
-
-or download a pretrained one from [2020-09-23T17-56-33_imagenet_vqgan](https://k00.fr/u0j2dtac)
-and place under `logs`. If you trained your own, adjust the path in the config
-key `model.params.first_stage_config.params.ckpt_path` of
-`configs/drin_transformer.yaml`.
-
-Train a VQGAN on Depth Maps of ImageNet with
-```
-python main.py --base configs/imagenetdepth_vqgan.yaml -t True --gpus 0,
-```
-
-or download a pretrained one from [2020-11-03T15-34-24_imagenetdepth_vqgan](https://k00.fr/55rlxs6i)
-and place under `logs`. If you trained your own, adjust the path in the config
-key `model.params.cond_stage_config.params.ckpt_path` of
-`configs/drin_transformer.yaml`.
-
-To train the transformer, run
-```
-python main.py --base configs/drin_transformer.yaml -t True --gpus 0,
-```
-
-## More Resources
-### Comparing Different First Stage Models
-The reconstruction and compression capabilities of different fist stage models can be analyzed in this [colab notebook](https://colab.research.google.com/github/CompVis/taming-transformers/blob/master/scripts/reconstruction_usage.ipynb). 
-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's [DALL-E](https://github.com/openai/DALL-E) (which has f=8 and 8192 
-codebook entries).
-![firststages1](assets/first_stage_squirrels.png)
-![firststages2](assets/first_stage_mushrooms.png)
-
-### Other
-- A [video summary](https://www.youtube.com/watch?v=o7dqGcLDf0A&feature=emb_imp_woyt) by [Two Minute Papers](https://www.youtube.com/channel/UCbfYPyITQ-7l4upoX8nvctg).
-- A [video summary](https://www.youtube.com/watch?v=-wDSDtIAyWQ) by [Gradient Dude](https://www.youtube.com/c/GradientDude/about).
-- A [weights and biases report summarizing the paper](https://wandb.ai/ayush-thakur/taming-transformer/reports/-Overview-Taming-Transformers-for-High-Resolution-Image-Synthesis---Vmlldzo0NjEyMTY)
-by [ayulockin](https://github.com/ayulockin).
-- A [video summary](https://www.youtube.com/watch?v=JfUTd8fjtX8&feature=emb_imp_woyt) by [What's AI](https://www.youtube.com/channel/UCUzGQrN-lyyc0BWTYoJM_Sg).
-- Take a look at [ak9250's notebook](https://github.com/ak9250/taming-transformers/blob/master/tamingtransformerscolab.ipynb) if you want to run the streamlit demos on Colab.
-
-### Text-to-Image Optimization via CLIP
-VQGAN has been successfully used as an image generator guided by the [CLIP](https://github.com/openai/CLIP) model, both for pure image generation
-from scratch and image-to-image translation. We recommend the following notebooks/videos/resources:
-
- - [Advadnouns](https://twitter.com/advadnoun/status/1389316507134357506) Patreon and corresponding LatentVision notebooks: https://www.patreon.com/patronizeme
- - The [notebook]( https://colab.research.google.com/drive/1L8oL-vLJXVcRzCFbPwOoMkPKJ8-aYdPN) of [Rivers Have Wings](https://twitter.com/RiversHaveWings).
- - A [video](https://www.youtube.com/watch?v=90QDe6DQXF4&t=12s) explanation by [Dot CSV](https://www.youtube.com/channel/UCy5znSnfMsDwaLlROnZ7Qbg) (in Spanish, but English subtitles are available)
-
-![txt2img](assets/birddrawnbyachild.png)
-
-Text prompt: *'A bird drawn by a child'*
-
-## Shout-outs
-Thanks to everyone who makes their code and models available. In particular,
-
-- The architecture of our VQGAN is inspired by [Denoising Diffusion Probabilistic Models](https://github.com/hojonathanho/diffusion)
-- The very hackable transformer implementation [minGPT](https://github.com/karpathy/minGPT)
-- The good ol' [PatchGAN](https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix) and [Learned Perceptual Similarity (LPIPS)](https://github.com/richzhang/PerceptualSimilarity)
-
-## 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}
-}
-```

taming-transformers/configs/coco_cond_stage.yaml

@@ -1,49 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.vqgan.VQSegmentationModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    image_key: "segmentation"
-    n_labels: 183
-    ddconfig:
-      double_z: false
-      z_channels: 256
-      resolution: 256
-      in_channels: 183
-      out_ch: 183
-      ch: 128
-      ch_mult:
-      - 1
-      - 1
-      - 2
-      - 2
-      - 4
-      num_res_blocks: 2
-      attn_resolutions:
-      - 16
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.segmentation.BCELossWithQuant
-      params:
-        codebook_weight: 1.0
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 12
-    train:
-      target: taming.data.coco.CocoImagesAndCaptionsTrain
-      params:
-        size: 296
-        crop_size: 256
-        onehot_segmentation: true
-        use_stuffthing: true
-    validation:
-      target: taming.data.coco.CocoImagesAndCaptionsValidation
-      params:
-        size: 256
-        crop_size: 256
-        onehot_segmentation: true
-        use_stuffthing: true

taming-transformers/configs/coco_scene_images_transformer.yaml

@@ -1,80 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.cond_transformer.Net2NetTransformer
-  params:
-    cond_stage_key: objects_bbox
-    transformer_config:
-      target: taming.modules.transformer.mingpt.GPT
-      params:
-        vocab_size: 8192
-        block_size: 348  # = 256 + 92 = dim(vqgan_latent_space,16x16) + dim(conditional_builder.embedding_dim)
-        n_layer: 40
-        n_head: 16
-        n_embd: 1408
-        embd_pdrop: 0.1
-        resid_pdrop: 0.1
-        attn_pdrop: 0.1
-    first_stage_config:
-      target: taming.models.vqgan.VQModel
-      params:
-        ckpt_path: /path/to/coco_epoch117.ckpt  # https://heibox.uni-heidelberg.de/f/78dea9589974474c97c1/
-        embed_dim: 256
-        n_embed: 8192
-        ddconfig:
-          double_z: false
-          z_channels: 256
-          resolution: 256
-          in_channels: 3
-          out_ch: 3
-          ch: 128
-          ch_mult:
-          - 1
-          - 1
-          - 2
-          - 2
-          - 4
-          num_res_blocks: 2
-          attn_resolutions:
-          - 16
-          dropout: 0.0
-        lossconfig:
-          target: taming.modules.losses.DummyLoss
-    cond_stage_config:
-      target: taming.models.dummy_cond_stage.DummyCondStage
-      params:
-        conditional_key: objects_bbox
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 6
-    train:
-      target: taming.data.annotated_objects_coco.AnnotatedObjectsCoco
-      params:
-        data_path: data/coco_annotations_100  # substitute with path to full dataset
-        split: train
-        keys: [image, objects_bbox, file_name, annotations]
-        no_tokens: 8192
-        target_image_size: 256
-        min_object_area: 0.00001
-        min_objects_per_image: 2
-        max_objects_per_image: 30
-        crop_method: random-1d
-        random_flip: true
-        use_group_parameter: true
-        encode_crop: true
-    validation:
-      target: taming.data.annotated_objects_coco.AnnotatedObjectsCoco
-      params:
-        data_path: data/coco_annotations_100  # substitute with path to full dataset
-        split: validation
-        keys: [image, objects_bbox, file_name, annotations]
-        no_tokens: 8192
-        target_image_size: 256
-        min_object_area: 0.00001
-        min_objects_per_image: 2
-        max_objects_per_image: 30
-        crop_method: center
-        random_flip: false
-        use_group_parameter: true
-        encode_crop: true

taming-transformers/configs/custom_vqgan.yaml

@@ -1,43 +0,0 @@
-model:
-  base_learning_rate: 4.5e-6
-  target: taming.models.vqgan.VQModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    ddconfig:
-      double_z: False
-      z_channels: 256
-      resolution: 256
-      in_channels: 3
-      out_ch: 3
-      ch: 128
-      ch_mult: [ 1,1,2,2,4]  # num_down = len(ch_mult)-1
-      num_res_blocks: 2
-      attn_resolutions: [16]
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
-      params:
-        disc_conditional: False
-        disc_in_channels: 3
-        disc_start: 10000
-        disc_weight: 0.8
-        codebook_weight: 1.0
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 5
-    num_workers: 8
-    train:
-      target: taming.data.custom.CustomTrain
-      params:
-        training_images_list_file: some/training.txt
-        size: 256
-    validation:
-      target: taming.data.custom.CustomTest
-      params:
-        test_images_list_file: some/test.txt
-        size: 256
-

taming-transformers/configs/drin_transformer.yaml

@@ -1,77 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.cond_transformer.Net2NetTransformer
-  params:
-    cond_stage_key: depth
-    transformer_config:
-      target: taming.modules.transformer.mingpt.GPT
-      params:
-        vocab_size: 1024
-        block_size: 512
-        n_layer: 24
-        n_head: 16
-        n_embd: 1024
-    first_stage_config:
-      target: taming.models.vqgan.VQModel
-      params:
-        ckpt_path: logs/2020-09-23T17-56-33_imagenet_vqgan/checkpoints/last.ckpt
-        embed_dim: 256
-        n_embed: 1024
-        ddconfig:
-          double_z: false
-          z_channels: 256
-          resolution: 256
-          in_channels: 3
-          out_ch: 3
-          ch: 128
-          ch_mult:
-          - 1
-          - 1
-          - 2
-          - 2
-          - 4
-          num_res_blocks: 2
-          attn_resolutions:
-          - 16
-          dropout: 0.0
-        lossconfig:
-          target: taming.modules.losses.DummyLoss
-    cond_stage_config:
-      target: taming.models.vqgan.VQModel
-      params:
-        ckpt_path: logs/2020-11-03T15-34-24_imagenetdepth_vqgan/checkpoints/last.ckpt
-        embed_dim: 256
-        n_embed: 1024
-        ddconfig:
-          double_z: false
-          z_channels: 256
-          resolution: 256
-          in_channels: 1
-          out_ch: 1
-          ch: 128
-          ch_mult:
-          - 1
-          - 1
-          - 2
-          - 2
-          - 4
-          num_res_blocks: 2
-          attn_resolutions:
-          - 16
-          dropout: 0.0
-        lossconfig:
-          target: taming.modules.losses.DummyLoss
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 2
-    num_workers: 8
-    train:
-      target: taming.data.imagenet.RINTrainWithDepth
-      params:
-        size: 256
-    validation:
-      target: taming.data.imagenet.RINValidationWithDepth
-      params:
-        size: 256

taming-transformers/configs/faceshq_transformer.yaml

@@ -1,61 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.cond_transformer.Net2NetTransformer
-  params:
-    cond_stage_key: coord
-    transformer_config:
-      target: taming.modules.transformer.mingpt.GPT
-      params:
-        vocab_size: 1024
-        block_size: 512
-        n_layer: 24
-        n_head: 16
-        n_embd: 1024
-    first_stage_config:
-      target: taming.models.vqgan.VQModel
-      params:
-        ckpt_path: logs/2020-11-09T13-33-36_faceshq_vqgan/checkpoints/last.ckpt
-        embed_dim: 256
-        n_embed: 1024
-        ddconfig:
-          double_z: false
-          z_channels: 256
-          resolution: 256
-          in_channels: 3
-          out_ch: 3
-          ch: 128
-          ch_mult:
-          - 1
-          - 1
-          - 2
-          - 2
-          - 4
-          num_res_blocks: 2
-          attn_resolutions:
-          - 16
-          dropout: 0.0
-        lossconfig:
-          target: taming.modules.losses.DummyLoss
-    cond_stage_config:
-      target: taming.modules.misc.coord.CoordStage
-      params:
-        n_embed: 1024
-        down_factor: 16
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 2
-    num_workers: 8
-    train:
-      target: taming.data.faceshq.FacesHQTrain
-      params:
-        size: 256
-        crop_size: 256
-        coord: True
-    validation:
-      target: taming.data.faceshq.FacesHQValidation
-      params:
-        size: 256
-        crop_size: 256
-        coord: True

taming-transformers/configs/faceshq_vqgan.yaml

@@ -1,42 +0,0 @@
-model:
-  base_learning_rate: 4.5e-6
-  target: taming.models.vqgan.VQModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    ddconfig:
-      double_z: False
-      z_channels: 256
-      resolution: 256
-      in_channels: 3
-      out_ch: 3
-      ch: 128
-      ch_mult: [ 1,1,2,2,4]  # num_down = len(ch_mult)-1
-      num_res_blocks: 2
-      attn_resolutions: [16]
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
-      params:
-        disc_conditional: False
-        disc_in_channels: 3
-        disc_start: 30001
-        disc_weight: 0.8
-        codebook_weight: 1.0
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 3
-    num_workers: 8
-    train:
-      target: taming.data.faceshq.FacesHQTrain
-      params:
-        size: 256
-        crop_size: 256
-    validation:
-      target: taming.data.faceshq.FacesHQValidation
-      params:
-        size: 256
-        crop_size: 256

taming-transformers/configs/imagenet_vqgan.yaml

@@ -1,42 +0,0 @@
-model:
-  base_learning_rate: 4.5e-6
-  target: taming.models.vqgan.VQModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    ddconfig:
-      double_z: False
-      z_channels: 256
-      resolution: 256
-      in_channels: 3
-      out_ch: 3
-      ch: 128
-      ch_mult: [ 1,1,2,2,4]  # num_down = len(ch_mult)-1
-      num_res_blocks: 2
-      attn_resolutions: [16]
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
-      params:
-        disc_conditional: False
-        disc_in_channels: 3
-        disc_start: 250001
-        disc_weight: 0.8
-        codebook_weight: 1.0
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 12
-    num_workers: 24
-    train:
-      target: taming.data.imagenet.ImageNetTrain
-      params:
-        config:
-          size: 256
-    validation:
-      target: taming.data.imagenet.ImageNetValidation
-      params:
-        config:
-          size: 256

taming-transformers/configs/imagenetdepth_vqgan.yaml

@@ -1,41 +0,0 @@
-model:
-  base_learning_rate: 4.5e-6
-  target: taming.models.vqgan.VQModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    image_key: depth
-    ddconfig:
-      double_z: False
-      z_channels: 256
-      resolution: 256
-      in_channels: 1
-      out_ch: 1
-      ch: 128
-      ch_mult: [ 1,1,2,2,4]  # num_down = len(ch_mult)-1
-      num_res_blocks: 2
-      attn_resolutions: [16]
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.vqperceptual.VQLPIPSWithDiscriminator
-      params:
-        disc_conditional: False
-        disc_in_channels: 1
-        disc_start: 50001
-        disc_weight: 0.75
-        codebook_weight: 1.0
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 3
-    num_workers: 8
-    train:
-      target: taming.data.imagenet.ImageNetTrainWithDepth
-      params:
-        size: 256
-    validation:
-      target: taming.data.imagenet.ImageNetValidationWithDepth
-      params:
-        size: 256

taming-transformers/configs/open_images_scene_images_transformer.yaml

@@ -1,86 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.cond_transformer.Net2NetTransformer
-  params:
-    cond_stage_key: objects_bbox
-    transformer_config:
-      target: taming.modules.transformer.mingpt.GPT
-      params:
-        vocab_size: 8192
-        block_size: 348  # = 256 + 92 = dim(vqgan_latent_space,16x16) + dim(conditional_builder.embedding_dim)
-        n_layer: 36
-        n_head: 16
-        n_embd: 1536
-        embd_pdrop: 0.1
-        resid_pdrop: 0.1
-        attn_pdrop: 0.1
-    first_stage_config:
-      target: taming.models.vqgan.VQModel
-      params:
-        ckpt_path: /path/to/coco_oi_epoch12.ckpt  # https://heibox.uni-heidelberg.de/f/461d9a9f4fcf48ab84f4/
-        embed_dim: 256
-        n_embed: 8192
-        ddconfig:
-          double_z: false
-          z_channels: 256
-          resolution: 256
-          in_channels: 3
-          out_ch: 3
-          ch: 128
-          ch_mult:
-          - 1
-          - 1
-          - 2
-          - 2
-          - 4
-          num_res_blocks: 2
-          attn_resolutions:
-          - 16
-          dropout: 0.0
-        lossconfig:
-          target: taming.modules.losses.DummyLoss
-    cond_stage_config:
-      target: taming.models.dummy_cond_stage.DummyCondStage
-      params:
-        conditional_key: objects_bbox
-
-data:
-  target: main.DataModuleFromConfig
-  params:
-    batch_size: 6
-    train:
-      target: taming.data.annotated_objects_open_images.AnnotatedObjectsOpenImages
-      params:
-        data_path: data/open_images_annotations_100  # substitute with path to full dataset
-        split: train
-        keys: [image, objects_bbox, file_name, annotations]
-        no_tokens: 8192
-        target_image_size: 256
-        category_allow_list_target: taming.data.open_images_helper.top_300_classes_plus_coco_compatibility
-        category_mapping_target: taming.data.open_images_helper.open_images_unify_categories_for_coco
-        min_object_area: 0.0001
-        min_objects_per_image: 2
-        max_objects_per_image: 30
-        crop_method: random-2d
-        random_flip: true
-        use_group_parameter: true
-        use_additional_parameters: true
-        encode_crop: true
-    validation:
-      target: taming.data.annotated_objects_open_images.AnnotatedObjectsOpenImages
-      params:
-        data_path: data/open_images_annotations_100  # substitute with path to full dataset
-        split: validation
-        keys: [image, objects_bbox, file_name, annotations]
-        no_tokens: 8192
-        target_image_size: 256
-        category_allow_list_target: taming.data.open_images_helper.top_300_classes_plus_coco_compatibility
-        category_mapping_target: taming.data.open_images_helper.open_images_unify_categories_for_coco
-        min_object_area: 0.0001
-        min_objects_per_image: 2
-        max_objects_per_image: 30
-        crop_method: center
-        random_flip: false
-        use_group_parameter: true
-        use_additional_parameters: true
-        encode_crop: true

taming-transformers/configs/sflckr_cond_stage.yaml

@@ -1,43 +0,0 @@
-model:
-  base_learning_rate: 4.5e-06
-  target: taming.models.vqgan.VQSegmentationModel
-  params:
-    embed_dim: 256
-    n_embed: 1024
-    image_key: "segmentation"
-    n_labels: 182
-    ddconfig:
-      double_z: false
-      z_channels: 256
-      resolution: 256
-      in_channels: 182
-      out_ch: 182
-      ch: 128
-      ch_mult:
-      - 1
-      - 1
-      - 2
-      - 2
-      - 4
-      num_res_blocks: 2
-      attn_resolutions:
-      - 16
-      dropout: 0.0
-
-    lossconfig:
-      target: taming.modules.losses.segmentation.BCELossWithQuant
-      params:
-        codebook_weight: 1.0
-
-data:
-  target: cutlit.DataModuleFromConfig
-  params:
-    batch_size: 12
-    train:
-      target: taming.data.sflckr.Examples # adjust
-      params:
-        size: 256
-    validation:
-      target: taming.data.sflckr.Examples # adjust
-      params:
-        size: 256

taming-transformers/environment.yaml

@@ -1,25 +0,0 @@
-name: taming
-channels:
-  - pytorch
-  - defaults
-dependencies:
-  - python=3.8.5
-  - pip=20.3
-  - cudatoolkit=10.2
-  - pytorch=1.7.0
-  - torchvision=0.8.1
-  - numpy=1.19.2
-  - pip:
-    - albumentations==0.4.3
-    - opencv-python==4.1.2.30
-    - pudb==2019.2
-    - imageio==2.9.0
-    - imageio-ffmpeg==0.4.2
-    - pytorch-lightning==1.0.8
-    - omegaconf==2.0.0
-    - test-tube>=0.7.5
-    - streamlit>=0.73.1
-    - einops==0.3.0
-    - more-itertools>=8.0.0
-    - transformers==4.3.1
-    - -e .

taming-transformers/main.py

@@ -1,585 +0,0 @@
-import argparse, os, sys, datetime, glob, importlib
-from omegaconf import OmegaConf
-import numpy as np
-from PIL import Image
-import torch
-import torchvision
-from torch.utils.data import random_split, DataLoader, Dataset
-import pytorch_lightning as pl
-from pytorch_lightning import seed_everything
-from pytorch_lightning.trainer import Trainer
-from pytorch_lightning.callbacks import ModelCheckpoint, Callback, LearningRateMonitor
-from pytorch_lightning.utilities import rank_zero_only
-
-from taming.data.utils import custom_collate
-
-
-def get_obj_from_str(string, reload=False):
-    module, cls = string.rsplit(".", 1)
-    if reload:
-        module_imp = importlib.import_module(module)
-        importlib.reload(module_imp)
-    return getattr(importlib.import_module(module, package=None), cls)
-
-
-def get_parser(**parser_kwargs):
-    def str2bool(v):
-        if isinstance(v, bool):
-            return v
-        if v.lower() in ("yes", "true", "t", "y", "1"):
-            return True
-        elif v.lower() in ("no", "false", "f", "n", "0"):
-            return False
-        else:
-            raise argparse.ArgumentTypeError("Boolean value expected.")
-
-    parser = argparse.ArgumentParser(**parser_kwargs)
-    parser.add_argument(
-        "-n",
-        "--name",
-        type=str,
-        const=True,
-        default="",
-        nargs="?",
-        help="postfix for logdir",
-    )
-    parser.add_argument(
-        "-r",
-        "--resume",
-        type=str,
-        const=True,
-        default="",
-        nargs="?",
-        help="resume from logdir or checkpoint in logdir",
-    )
-    parser.add_argument(
-        "-b",
-        "--base",
-        nargs="*",
-        metavar="base_config.yaml",
-        help="paths to base configs. Loaded from left-to-right. "
-        "Parameters can be overwritten or added with command-line options of the form `--key value`.",
-        default=list(),
-    )
-    parser.add_argument(
-        "-t",
-        "--train",
-        type=str2bool,
-        const=True,
-        default=False,
-        nargs="?",
-        help="train",
-    )
-    parser.add_argument(
-        "--no-test",
-        type=str2bool,
-        const=True,
-        default=False,
-        nargs="?",
-        help="disable test",
-    )
-    parser.add_argument("-p", "--project", help="name of new or path to existing project")
-    parser.add_argument(
-        "-d",
-        "--debug",
-        type=str2bool,
-        nargs="?",
-        const=True,
-        default=False,
-        help="enable post-mortem debugging",
-    )
-    parser.add_argument(
-        "-s",
-        "--seed",
-        type=int,
-        default=23,
-        help="seed for seed_everything",
-    )
-    parser.add_argument(
-        "-f",
-        "--postfix",
-        type=str,
-        default="",
-        help="post-postfix for default name",
-    )
-
-    return parser
-
-
-def nondefault_trainer_args(opt):
-    parser = argparse.ArgumentParser()
-    parser = Trainer.add_argparse_args(parser)
-    args = parser.parse_args([])
-    return sorted(k for k in vars(args) if getattr(opt, k) != getattr(args, k))
-
-
-def instantiate_from_config(config):
-    if not "target" in config:
-        raise KeyError("Expected key `target` to instantiate.")
-    return get_obj_from_str(config["target"])(**config.get("params", dict()))
-
-
-class WrappedDataset(Dataset):
-    """Wraps an arbitrary object with __len__ and __getitem__ into a pytorch dataset"""
-    def __init__(self, dataset):
-        self.data = dataset
-
-    def __len__(self):
-        return len(self.data)
-
-    def __getitem__(self, idx):
-        return self.data[idx]
-
-
-class DataModuleFromConfig(pl.LightningDataModule):
-    def __init__(self, batch_size, train=None, validation=None, test=None,
-                 wrap=False, num_workers=None):
-        super().__init__()
-        self.batch_size = batch_size
-        self.dataset_configs = dict()
-        self.num_workers = num_workers if num_workers is not None else batch_size*2
-        if train is not None:
-            self.dataset_configs["train"] = train
-            self.train_dataloader = self._train_dataloader
-        if validation is not None:
-            self.dataset_configs["validation"] = validation
-            self.val_dataloader = self._val_dataloader
-        if test is not None:
-            self.dataset_configs["test"] = test
-            self.test_dataloader = self._test_dataloader
-        self.wrap = wrap
-
-    def prepare_data(self):
-        for data_cfg in self.dataset_configs.values():
-            instantiate_from_config(data_cfg)
-
-    def setup(self, stage=None):
-        self.datasets = dict(
-            (k, instantiate_from_config(self.dataset_configs[k]))
-            for k in self.dataset_configs)
-        if self.wrap:
-            for k in self.datasets:
-                self.datasets[k] = WrappedDataset(self.datasets[k])
-
-    def _train_dataloader(self):
-        return DataLoader(self.datasets["train"], batch_size=self.batch_size,
-                          num_workers=self.num_workers, shuffle=True, collate_fn=custom_collate)
-
-    def _val_dataloader(self):
-        return DataLoader(self.datasets["validation"],
-                          batch_size=self.batch_size,
-                          num_workers=self.num_workers, collate_fn=custom_collate)
-
-    def _test_dataloader(self):
-        return DataLoader(self.datasets["test"], batch_size=self.batch_size,
-                          num_workers=self.num_workers, collate_fn=custom_collate)
-
-
-class SetupCallback(Callback):
-    def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config):
-        super().__init__()
-        self.resume = resume
-        self.now = now
-        self.logdir = logdir
-        self.ckptdir = ckptdir
-        self.cfgdir = cfgdir
-        self.config = config
-        self.lightning_config = lightning_config
-
-    def on_pretrain_routine_start(self, trainer, pl_module):
-        if trainer.global_rank == 0:
-            # Create logdirs and save configs
-            os.makedirs(self.logdir, exist_ok=True)
-            os.makedirs(self.ckptdir, exist_ok=True)
-            os.makedirs(self.cfgdir, exist_ok=True)
-
-            print("Project config")
-            print(self.config.pretty())
-            OmegaConf.save(self.config,
-                           os.path.join(self.cfgdir, "{}-project.yaml".format(self.now)))
-
-            print("Lightning config")
-            print(self.lightning_config.pretty())
-            OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}),
-                           os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now)))
-
-        else:
-            # ModelCheckpoint callback created log directory --- remove it
-            if not self.resume and os.path.exists(self.logdir):
-                dst, name = os.path.split(self.logdir)
-                dst = os.path.join(dst, "child_runs", name)
-                os.makedirs(os.path.split(dst)[0], exist_ok=True)
-                try:
-                    os.rename(self.logdir, dst)
-                except FileNotFoundError:
-                    pass
-
-
-class ImageLogger(Callback):
-    def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True):
-        super().__init__()
-        self.batch_freq = batch_frequency
-        self.max_images = max_images
-        self.logger_log_images = {
-            pl.loggers.WandbLogger: self._wandb,
-            pl.loggers.TestTubeLogger: self._testtube,
-        }
-        self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)]
-        if not increase_log_steps:
-            self.log_steps = [self.batch_freq]
-        self.clamp = clamp
-
-    @rank_zero_only
-    def _wandb(self, pl_module, images, batch_idx, split):
-        raise ValueError("No way wandb")
-        grids = dict()
-        for k in images:
-            grid = torchvision.utils.make_grid(images[k])
-            grids[f"{split}/{k}"] = wandb.Image(grid)
-        pl_module.logger.experiment.log(grids)
-
-    @rank_zero_only
-    def _testtube(self, pl_module, images, batch_idx, split):
-        for k in images:
-            grid = torchvision.utils.make_grid(images[k])
-            grid = (grid+1.0)/2.0 # -1,1 -> 0,1; c,h,w
-
-            tag = f"{split}/{k}"
-            pl_module.logger.experiment.add_image(
-                tag, grid,
-                global_step=pl_module.global_step)
-
-    @rank_zero_only
-    def log_local(self, save_dir, split, images,
-                  global_step, current_epoch, batch_idx):
-        root = os.path.join(save_dir, "images", split)
-        for k in images:
-            grid = torchvision.utils.make_grid(images[k], nrow=4)
-
-            grid = (grid+1.0)/2.0 # -1,1 -> 0,1; c,h,w
-            grid = grid.transpose(0,1).transpose(1,2).squeeze(-1)
-            grid = grid.numpy()
-            grid = (grid*255).astype(np.uint8)
-            filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(
-                k,
-                global_step,
-                current_epoch,
-                batch_idx)
-            path = os.path.join(root, filename)
-            os.makedirs(os.path.split(path)[0], exist_ok=True)
-            Image.fromarray(grid).save(path)
-
-    def log_img(self, pl_module, batch, batch_idx, split="train"):
-        if (self.check_frequency(batch_idx) and  # batch_idx % self.batch_freq == 0
-                hasattr(pl_module, "log_images") and
-                callable(pl_module.log_images) and
-                self.max_images > 0):
-            logger = type(pl_module.logger)
-
-            is_train = pl_module.training
-            if is_train:
-                pl_module.eval()
-
-            with torch.no_grad():
-                images = pl_module.log_images(batch, split=split, pl_module=pl_module)
-
-            for k in images:
-                N = min(images[k].shape[0], self.max_images)
-                images[k] = images[k][:N]
-                if isinstance(images[k], torch.Tensor):
-                    images[k] = images[k].detach().cpu()
-                    if self.clamp:
-                        images[k] = torch.clamp(images[k], -1., 1.)
-
-            self.log_local(pl_module.logger.save_dir, split, images,
-                           pl_module.global_step, pl_module.current_epoch, batch_idx)
-
-            logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None)
-            logger_log_images(pl_module, images, pl_module.global_step, split)
-
-            if is_train:
-                pl_module.train()
-
-    def check_frequency(self, batch_idx):
-        if (batch_idx % self.batch_freq) == 0 or (batch_idx in self.log_steps):
-            try:
-                self.log_steps.pop(0)
-            except IndexError:
-                pass
-            return True
-        return False
-
-    def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
-        self.log_img(pl_module, batch, batch_idx, split="train")
-
-    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
-        self.log_img(pl_module, batch, batch_idx, split="val")
-
-
-
-if __name__ == "__main__":
-    # custom parser to specify config files, train, test and debug mode,
-    # postfix, resume.
-    # `--key value` arguments are interpreted as arguments to the trainer.
-    # `nested.key=value` arguments are interpreted as config parameters.
-    # configs are merged from left-to-right followed by command line parameters.
-
-    # model:
-    #   base_learning_rate: float
-    #   target: path to lightning module
-    #   params:
-    #       key: value
-    # data:
-    #   target: main.DataModuleFromConfig
-    #   params:
-    #      batch_size: int
-    #      wrap: bool
-    #      train:
-    #          target: path to train dataset
-    #          params:
-    #              key: value
-    #      validation:
-    #          target: path to validation dataset
-    #          params:
-    #              key: value
-    #      test:
-    #          target: path to test dataset
-    #          params:
-    #              key: value
-    # lightning: (optional, has sane defaults and can be specified on cmdline)
-    #   trainer:
-    #       additional arguments to trainer
-    #   logger:
-    #       logger to instantiate
-    #   modelcheckpoint:
-    #       modelcheckpoint to instantiate
-    #   callbacks:
-    #       callback1:
-    #           target: importpath
-    #           params:
-    #               key: value
-
-    now = datetime.datetime.now().strftime("%Y-%m-%dT%H-%M-%S")
-
-    # add cwd for convenience and to make classes in this file available when
-    # running as `python main.py`
-    # (in particular `main.DataModuleFromConfig`)
-    sys.path.append(os.getcwd())
-
-    parser = get_parser()
-    parser = Trainer.add_argparse_args(parser)
-
-    opt, unknown = parser.parse_known_args()
-    if opt.name and opt.resume:
-        raise ValueError(
-            "-n/--name and -r/--resume cannot be specified both."
-            "If you want to resume training in a new log folder, "
-            "use -n/--name in combination with --resume_from_checkpoint"
-        )
-    if opt.resume:
-        if not os.path.exists(opt.resume):
-            raise ValueError("Cannot find {}".format(opt.resume))
-        if os.path.isfile(opt.resume):
-            paths = opt.resume.split("/")
-            idx = len(paths)-paths[::-1].index("logs")+1
-            logdir = "/".join(paths[:idx])
-            ckpt = opt.resume
-        else:
-            assert os.path.isdir(opt.resume), opt.resume
-            logdir = opt.resume.rstrip("/")
-            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
-
-        opt.resume_from_checkpoint = ckpt
-        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*.yaml")))
-        opt.base = base_configs+opt.base
-        _tmp = logdir.split("/")
-        nowname = _tmp[_tmp.index("logs")+1]
-    else:
-        if opt.name:
-            name = "_"+opt.name
-        elif opt.base:
-            cfg_fname = os.path.split(opt.base[0])[-1]
-            cfg_name = os.path.splitext(cfg_fname)[0]
-            name = "_"+cfg_name
-        else:
-            name = ""
-        nowname = now+name+opt.postfix
-        logdir = os.path.join("logs", nowname)
-
-    ckptdir = os.path.join(logdir, "checkpoints")
-    cfgdir = os.path.join(logdir, "configs")
-    seed_everything(opt.seed)
-
-    try:
-        # init and save configs
-        configs = [OmegaConf.load(cfg) for cfg in opt.base]
-        cli = OmegaConf.from_dotlist(unknown)
-        config = OmegaConf.merge(*configs, cli)
-        lightning_config = config.pop("lightning", OmegaConf.create())
-        # merge trainer cli with config
-        trainer_config = lightning_config.get("trainer", OmegaConf.create())
-        # default to ddp
-        trainer_config["distributed_backend"] = "ddp"
-        for k in nondefault_trainer_args(opt):
-            trainer_config[k] = getattr(opt, k)
-        if not "gpus" in trainer_config:
-            del trainer_config["distributed_backend"]
-            cpu = True
-        else:
-            gpuinfo = trainer_config["gpus"]
-            print(f"Running on GPUs {gpuinfo}")
-            cpu = False
-        trainer_opt = argparse.Namespace(**trainer_config)
-        lightning_config.trainer = trainer_config
-
-        # model
-        model = instantiate_from_config(config.model)
-
-        # trainer and callbacks
-        trainer_kwargs = dict()
-
-        # default logger configs
-        # NOTE wandb < 0.10.0 interferes with shutdown
-        # wandb >= 0.10.0 seems to fix it but still interferes with pudb
-        # debugging (wrongly sized pudb ui)
-        # thus prefer testtube for now
-        default_logger_cfgs = {
-            "wandb": {
-                "target": "pytorch_lightning.loggers.WandbLogger",
-                "params": {
-                    "name": nowname,
-                    "save_dir": logdir,
-                    "offline": opt.debug,
-                    "id": nowname,
-                }
-            },
-            "testtube": {
-                "target": "pytorch_lightning.loggers.TestTubeLogger",
-                "params": {
-                    "name": "testtube",
-                    "save_dir": logdir,
-                }
-            },
-        }
-        default_logger_cfg = default_logger_cfgs["testtube"]
-        logger_cfg = lightning_config.logger or OmegaConf.create()
-        logger_cfg = OmegaConf.merge(default_logger_cfg, logger_cfg)
-        trainer_kwargs["logger"] = instantiate_from_config(logger_cfg)
-
-        # modelcheckpoint - use TrainResult/EvalResult(checkpoint_on=metric) to
-        # specify which metric is used to determine best models
-        default_modelckpt_cfg = {
-            "target": "pytorch_lightning.callbacks.ModelCheckpoint",
-            "params": {
-                "dirpath": ckptdir,
-                "filename": "{epoch:06}",
-                "verbose": True,
-                "save_last": True,
-            }
-        }
-        if hasattr(model, "monitor"):
-            print(f"Monitoring {model.monitor} as checkpoint metric.")
-            default_modelckpt_cfg["params"]["monitor"] = model.monitor
-            default_modelckpt_cfg["params"]["save_top_k"] = 3
-
-        modelckpt_cfg = lightning_config.modelcheckpoint or OmegaConf.create()
-        modelckpt_cfg = OmegaConf.merge(default_modelckpt_cfg, modelckpt_cfg)
-        trainer_kwargs["checkpoint_callback"] = instantiate_from_config(modelckpt_cfg)
-
-        # add callback which sets up log directory
-        default_callbacks_cfg = {
-            "setup_callback": {
-                "target": "main.SetupCallback",
-                "params": {
-                    "resume": opt.resume,
-                    "now": now,
-                    "logdir": logdir,
-                    "ckptdir": ckptdir,
-                    "cfgdir": cfgdir,
-                    "config": config,
-                    "lightning_config": lightning_config,
-                }
-            },
-            "image_logger": {
-                "target": "main.ImageLogger",
-                "params": {
-                    "batch_frequency": 750,
-                    "max_images": 4,
-                    "clamp": True
-                }
-            },
-            "learning_rate_logger": {
-                "target": "main.LearningRateMonitor",
-                "params": {
-                    "logging_interval": "step",
-                    #"log_momentum": True
-                }
-            },
-        }
-        callbacks_cfg = lightning_config.callbacks or OmegaConf.create()
-        callbacks_cfg = OmegaConf.merge(default_callbacks_cfg, callbacks_cfg)
-        trainer_kwargs["callbacks"] = [instantiate_from_config(callbacks_cfg[k]) for k in callbacks_cfg]
-
-        trainer = Trainer.from_argparse_args(trainer_opt, **trainer_kwargs)
-
-        # data
-        data = instantiate_from_config(config.data)
-        # NOTE according to https://pytorch-lightning.readthedocs.io/en/latest/datamodules.html
-        # calling these ourselves should not be necessary but it is.
-        # lightning still takes care of proper multiprocessing though
-        data.prepare_data()
-        data.setup()
-
-        # configure learning rate
-        bs, base_lr = config.data.params.batch_size, config.model.base_learning_rate
-        if not cpu:
-            ngpu = len(lightning_config.trainer.gpus.strip(",").split(','))
-        else:
-            ngpu = 1
-        accumulate_grad_batches = lightning_config.trainer.accumulate_grad_batches or 1
-        print(f"accumulate_grad_batches = {accumulate_grad_batches}")
-        lightning_config.trainer.accumulate_grad_batches = accumulate_grad_batches
-        model.learning_rate = accumulate_grad_batches * ngpu * bs * base_lr
-        print("Setting learning rate to {:.2e} = {} (accumulate_grad_batches) * {} (num_gpus) * {} (batchsize) * {:.2e} (base_lr)".format(
-            model.learning_rate, accumulate_grad_batches, ngpu, bs, base_lr))
-
-        # allow checkpointing via USR1
-        def melk(*args, **kwargs):
-            # run all checkpoint hooks
-            if trainer.global_rank == 0:
-                print("Summoning checkpoint.")
-                ckpt_path = os.path.join(ckptdir, "last.ckpt")
-                trainer.save_checkpoint(ckpt_path)
-
-        def divein(*args, **kwargs):
-            if trainer.global_rank == 0:
-                import pudb; pudb.set_trace()
-
-        import signal
-        signal.signal(signal.SIGUSR1, melk)
-        signal.signal(signal.SIGUSR2, divein)
-
-        # run
-        if opt.train:
-            try:
-                trainer.fit(model, data)
-            except Exception:
-                melk()
-                raise
-        if not opt.no_test and not trainer.interrupted:
-            trainer.test(model, data)
-    except Exception:
-        if opt.debug and trainer.global_rank==0:
-            try:
-                import pudb as debugger
-            except ImportError:
-                import pdb as debugger
-            debugger.post_mortem()
-        raise
-    finally:
-        # move newly created debug project to debug_runs
-        if opt.debug and not opt.resume and trainer.global_rank==0:
-            dst, name = os.path.split(logdir)
-            dst = os.path.join(dst, "debug_runs", name)
-            os.makedirs(os.path.split(dst)[0], exist_ok=True)
-            os.rename(logdir, dst)

taming-transformers/scripts/extract_depth.py

@@ -1,112 +0,0 @@
-import os
-import torch
-import numpy as np
-from tqdm import trange
-from PIL import Image
-
-
-def get_state(gpu):
-    import torch
-    midas = torch.hub.load("intel-isl/MiDaS", "MiDaS")
-    if gpu:
-        midas.cuda()
-    midas.eval()
-
-    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
-    transform = midas_transforms.default_transform
-
-    state = {"model": midas,
-             "transform": transform}
-    return state
-
-
-def depth_to_rgba(x):
-    assert x.dtype == np.float32
-    assert len(x.shape) == 2
-    y = x.copy()
-    y.dtype = np.uint8
-    y = y.reshape(x.shape+(4,))
-    return np.ascontiguousarray(y)
-
-
-def rgba_to_depth(x):
-    assert x.dtype == np.uint8
-    assert len(x.shape) == 3 and x.shape[2] == 4
-    y = x.copy()
-    y.dtype = np.float32
-    y = y.reshape(x.shape[:2])
-    return np.ascontiguousarray(y)
-
-
-def run(x, state):
-    model = state["model"]
-    transform = state["transform"]
-    hw = x.shape[:2]
-    with torch.no_grad():
-        prediction = model(transform((x + 1.0) * 127.5).cuda())
-        prediction = torch.nn.functional.interpolate(
-            prediction.unsqueeze(1),
-            size=hw,
-            mode="bicubic",
-            align_corners=False,
-        ).squeeze()
-        output = prediction.cpu().numpy()
-    return output
-
-
-def get_filename(relpath, level=-2):
-    # save class folder structure and filename:
-    fn = relpath.split(os.sep)[level:]
-    folder = fn[-2]
-    file   = fn[-1].split('.')[0]
-    return folder, file
-
-
-def save_depth(dataset, path, debug=False):
-    os.makedirs(path)
-    N = len(dset)
-    if debug:
-        N = 10
-    state = get_state(gpu=True)
-    for idx in trange(N, desc="Data"):
-        ex = dataset[idx]
-        image, relpath = ex["image"], ex["relpath"]
-        folder, filename = get_filename(relpath)
-        # prepare
-        folderabspath = os.path.join(path, folder)
-        os.makedirs(folderabspath, exist_ok=True)
-        savepath = os.path.join(folderabspath, filename)
-        # run model
-        xout = run(image, state)
-        I = depth_to_rgba(xout)
-        Image.fromarray(I).save("{}.png".format(savepath))
-
-
-if __name__ == "__main__":
-    from taming.data.imagenet import ImageNetTrain, ImageNetValidation
-    out = "data/imagenet_depth"
-    if not os.path.exists(out):
-        print("Please create a folder or symlink '{}' to extract depth data ".format(out) +
-              "(be prepared that the output size will be larger than ImageNet itself).")
-        exit(1)
-
-    # go
-    dset = ImageNetValidation()
-    abspath = os.path.join(out, "val")
-    if os.path.exists(abspath):
-        print("{} exists - not doing anything.".format(abspath))
-    else:
-        print("preparing {}".format(abspath))
-        save_depth(dset, abspath)
-        print("done with validation split")
-
-    dset = ImageNetTrain()
-    abspath = os.path.join(out, "train")
-    if os.path.exists(abspath):
-        print("{} exists - not doing anything.".format(abspath))
-    else:
-        print("preparing {}".format(abspath))
-        save_depth(dset, abspath)
-        print("done with train split")
-
-    print("done done.")

taming-transformers/scripts/extract_segmentation.py

@@ -1,130 +0,0 @@
-import sys, os
-import numpy as np
-import scipy
-import torch
-import torch.nn as nn
-from scipy import ndimage
-from tqdm import tqdm, trange
-from PIL import Image
-import torch.hub
-import torchvision
-import torch.nn.functional as F
-
-# download deeplabv2_resnet101_msc-cocostuff164k-100000.pth from
-# https://github.com/kazuto1011/deeplab-pytorch/releases/download/v1.0/deeplabv2_resnet101_msc-cocostuff164k-100000.pth
-# and put the path here
-CKPT_PATH = "TODO"
-
-rescale = lambda x: (x + 1.) / 2.
-
-def rescale_bgr(x):
-    x = (x+1)*127.5
-    x = torch.flip(x, dims=[0])
-    return x
-
-
-class COCOStuffSegmenter(nn.Module):
-    def __init__(self, config):
-        super().__init__()
-        self.config = config
-        self.n_labels = 182
-        model = torch.hub.load("kazuto1011/deeplab-pytorch", "deeplabv2_resnet101", n_classes=self.n_labels)
-        ckpt_path = CKPT_PATH
-        model.load_state_dict(torch.load(ckpt_path))
-        self.model = model
-
-        normalize = torchvision.transforms.Normalize(mean=self.mean, std=self.std)
-        self.image_transform = torchvision.transforms.Compose([
-            torchvision.transforms.Lambda(lambda image: torch.stack(
-                [normalize(rescale_bgr(x)) for x in image]))
-        ])
-
-    def forward(self, x, upsample=None):
-        x = self._pre_process(x)
-        x = self.model(x)
-        if upsample is not None:
-            x = torch.nn.functional.upsample_bilinear(x, size=upsample)
-        return x
-
-    def _pre_process(self, x):
-        x = self.image_transform(x)
-        return x
-
-    @property
-    def mean(self):
-        # bgr
-        return [104.008, 116.669, 122.675]
-
-    @property
-    def std(self):
-        return [1.0, 1.0, 1.0]
-
-    @property
-    def input_size(self):
-        return [3, 224, 224]
-
-
-def run_model(img, model):
-    model = model.eval()
-    with torch.no_grad():
-        segmentation = model(img, upsample=(img.shape[2], img.shape[3]))
-        segmentation = torch.argmax(segmentation, dim=1, keepdim=True)
-    return segmentation.detach().cpu()
-
-
-def get_input(batch, k):
-    x = batch[k]
-    if len(x.shape) == 3:
-        x = x[..., None]
-    x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
-    return x.float()
-
-
-def save_segmentation(segmentation, path):
-    # --> class label to uint8, save as png
-    os.makedirs(os.path.dirname(path), exist_ok=True)
-    assert len(segmentation.shape)==4
-    assert segmentation.shape[0]==1
-    for seg in segmentation:
-        seg = seg.permute(1,2,0).numpy().squeeze().astype(np.uint8)
-        seg = Image.fromarray(seg)
-        seg.save(path)
-
-
-def iterate_dataset(dataloader, destpath, model):
-    os.makedirs(destpath, exist_ok=True)
-    num_processed = 0
-    for i, batch in tqdm(enumerate(dataloader), desc="Data"):
-        try:
-            img = get_input(batch, "image")
-            img = img.cuda()
-            seg = run_model(img, model)
-
-            path = batch["relative_file_path_"][0]
-            path = os.path.splitext(path)[0]
-
-            path = os.path.join(destpath, path + ".png")
-            save_segmentation(seg, path)
-            num_processed += 1
-        except Exception as e:
-            print(e)
-            print("but anyhow..")
-
-    print("Processed {} files. Bye.".format(num_processed))
-
-
-from taming.data.sflckr import Examples
-from torch.utils.data import DataLoader
-
-if __name__ == "__main__":
-    dest = sys.argv[1]
-    batchsize = 1
-    print("Running with batch-size {}, saving to {}...".format(batchsize, dest))
-
-    model = COCOStuffSegmenter({}).cuda()
-    print("Instantiated model.")
-
-    dataset = Examples()
-    dloader = DataLoader(dataset, batch_size=batchsize)
-    iterate_dataset(dataloader=dloader, destpath=dest, model=model)
-    print("done.")

taming-transformers/scripts/extract_submodel.py

@@ -1,17 +0,0 @@
-import torch
-import sys
-
-if __name__ == "__main__":
-    inpath = sys.argv[1]
-    outpath = sys.argv[2]
-    submodel = "cond_stage_model"
-    if len(sys.argv) > 3:
-        submodel = sys.argv[3]
-
-    print("Extracting {} from {} to {}.".format(submodel, inpath, outpath))
-
-    sd = torch.load(inpath, map_location="cpu")
-    new_sd = {"state_dict": dict((k.split(".", 1)[-1],v)
-                                 for k,v in sd["state_dict"].items()
-                                 if k.startswith("cond_stage_model"))}
-    torch.save(new_sd, outpath)

taming-transformers/scripts/make_samples.py

@@ -1,292 +0,0 @@
-import argparse, os, sys, glob, math, time
-import torch
-import numpy as np
-from omegaconf import OmegaConf
-from PIL import Image
-from main import instantiate_from_config, DataModuleFromConfig
-from torch.utils.data import DataLoader
-from torch.utils.data.dataloader import default_collate
-from tqdm import trange
-
-
-def save_image(x, path):
-    c,h,w = x.shape
-    assert c==3
-    x = ((x.detach().cpu().numpy().transpose(1,2,0)+1.0)*127.5).clip(0,255).astype(np.uint8)
-    Image.fromarray(x).save(path)
-
-
-@torch.no_grad()
-def run_conditional(model, dsets, outdir, top_k, temperature, batch_size=1):
-    if len(dsets.datasets) > 1:
-        split = sorted(dsets.datasets.keys())[0]
-        dset = dsets.datasets[split]
-    else:
-        dset = next(iter(dsets.datasets.values()))
-    print("Dataset: ", dset.__class__.__name__)
-    for start_idx in trange(0,len(dset)-batch_size+1,batch_size):
-        indices = list(range(start_idx, start_idx+batch_size))
-        example = default_collate([dset[i] for i in indices])
-
-        x = model.get_input("image", example).to(model.device)
-        for i in range(x.shape[0]):
-            save_image(x[i], os.path.join(outdir, "originals",
-                                          "{:06}.png".format(indices[i])))
-
-        cond_key = model.cond_stage_key
-        c = model.get_input(cond_key, example).to(model.device)
-
-        scale_factor = 1.0
-        quant_z, z_indices = model.encode_to_z(x)
-        quant_c, c_indices = model.encode_to_c(c)
-
-        cshape = quant_z.shape
-
-        xrec = model.first_stage_model.decode(quant_z)
-        for i in range(xrec.shape[0]):
-            save_image(xrec[i], os.path.join(outdir, "reconstructions",
-                                             "{:06}.png".format(indices[i])))
-
-        if cond_key == "segmentation":
-            # get image from segmentation mask
-            num_classes = c.shape[1]
-            c = torch.argmax(c, dim=1, keepdim=True)
-            c = torch.nn.functional.one_hot(c, num_classes=num_classes)
-            c = c.squeeze(1).permute(0, 3, 1, 2).float()
-            c = model.cond_stage_model.to_rgb(c)
-
-        idx = z_indices
-
-        half_sample = False
-        if half_sample:
-            start = idx.shape[1]//2
-        else:
-            start = 0
-
-        idx[:,start:] = 0
-        idx = idx.reshape(cshape[0],cshape[2],cshape[3])
-        start_i = start//cshape[3]
-        start_j = start %cshape[3]
-
-        cidx = c_indices
-        cidx = cidx.reshape(quant_c.shape[0],quant_c.shape[2],quant_c.shape[3])
-
-        sample = True
-
-        for i in range(start_i,cshape[2]-0):
-            if i <= 8:
-                local_i = i
-            elif cshape[2]-i < 8:
-                local_i = 16-(cshape[2]-i)
-            else:
-                local_i = 8
-            for j in range(start_j,cshape[3]-0):
-                if j <= 8:
-                    local_j = j
-                elif cshape[3]-j < 8:
-                    local_j = 16-(cshape[3]-j)
-                else:
-                    local_j = 8
-
-                i_start = i-local_i
-                i_end = i_start+16
-                j_start = j-local_j
-                j_end = j_start+16
-                patch = idx[:,i_start:i_end,j_start:j_end]
-                patch = patch.reshape(patch.shape[0],-1)
-                cpatch = cidx[:, i_start:i_end, j_start:j_end]
-                cpatch = cpatch.reshape(cpatch.shape[0], -1)
-                patch = torch.cat((cpatch, patch), dim=1)
-                logits,_ = model.transformer(patch[:,:-1])
-                logits = logits[:, -256:, :]
-                logits = logits.reshape(cshape[0],16,16,-1)
-                logits = logits[:,local_i,local_j,:]
-
-                logits = logits/temperature
-
-                if top_k is not None:
-                    logits = model.top_k_logits(logits, top_k)
-                # apply softmax to convert to probabilities
-                probs = torch.nn.functional.softmax(logits, dim=-1)
-                # sample from the distribution or take the most likely
-                if sample:
-                    ix = torch.multinomial(probs, num_samples=1)
-                else:
-                    _, ix = torch.topk(probs, k=1, dim=-1)
-                idx[:,i,j] = ix
-
-        xsample = model.decode_to_img(idx[:,:cshape[2],:cshape[3]], cshape)
-        for i in range(xsample.shape[0]):
-            save_image(xsample[i], os.path.join(outdir, "samples",
-                                                "{:06}.png".format(indices[i])))
-
-
-def get_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-r",
-        "--resume",
-        type=str,
-        nargs="?",
-        help="load from logdir or checkpoint in logdir",
-    )
-    parser.add_argument(
-        "-b",
-        "--base",
-        nargs="*",
-        metavar="base_config.yaml",
-        help="paths to base configs. Loaded from left-to-right. "
-        "Parameters can be overwritten or added with command-line options of the form `--key value`.",
-        default=list(),
-    )
-    parser.add_argument(
-        "-c",
-        "--config",
-        nargs="?",
-        metavar="single_config.yaml",
-        help="path to single config. If specified, base configs will be ignored "
-        "(except for the last one if left unspecified).",
-        const=True,
-        default="",
-    )
-    parser.add_argument(
-        "--ignore_base_data",
-        action="store_true",
-        help="Ignore data specification from base configs. Useful if you want "
-        "to specify a custom datasets on the command line.",
-    )
-    parser.add_argument(
-        "--outdir",
-        required=True,
-        type=str,
-        help="Where to write outputs to.",
-    )
-    parser.add_argument(
-        "--top_k",
-        type=int,
-        default=100,
-        help="Sample from among top-k predictions.",
-    )
-    parser.add_argument(
-        "--temperature",
-        type=float,
-        default=1.0,
-        help="Sampling temperature.",
-    )
-    return parser
-
-
-def load_model_from_config(config, sd, gpu=True, eval_mode=True):
-    if "ckpt_path" in config.params:
-        print("Deleting the restore-ckpt path from the config...")
-        config.params.ckpt_path = None
-    if "downsample_cond_size" in config.params:
-        print("Deleting downsample-cond-size from the config and setting factor=0.5 instead...")
-        config.params.downsample_cond_size = -1
-        config.params["downsample_cond_factor"] = 0.5
-    try:
-        if "ckpt_path" in config.params.first_stage_config.params:
-            config.params.first_stage_config.params.ckpt_path = None
-            print("Deleting the first-stage restore-ckpt path from the config...")
-        if "ckpt_path" in config.params.cond_stage_config.params:
-            config.params.cond_stage_config.params.ckpt_path = None
-            print("Deleting the cond-stage restore-ckpt path from the config...")
-    except:
-        pass
-
-    model = instantiate_from_config(config)
-    if sd is not None:
-        missing, unexpected = model.load_state_dict(sd, strict=False)
-        print(f"Missing Keys in State Dict: {missing}")
-        print(f"Unexpected Keys in State Dict: {unexpected}")
-    if gpu:
-        model.cuda()
-    if eval_mode:
-        model.eval()
-    return {"model": model}
-
-
-def get_data(config):
-    # get data
-    data = instantiate_from_config(config.data)
-    data.prepare_data()
-    data.setup()
-    return data
-
-
-def load_model_and_dset(config, ckpt, gpu, eval_mode):
-    # get data
-    dsets = get_data(config)   # calls data.config ...
-
-    # now load the specified checkpoint
-    if ckpt:
-        pl_sd = torch.load(ckpt, map_location="cpu")
-        global_step = pl_sd["global_step"]
-    else:
-        pl_sd = {"state_dict": None}
-        global_step = None
-    model = load_model_from_config(config.model,
-                                   pl_sd["state_dict"],
-                                   gpu=gpu,
-                                   eval_mode=eval_mode)["model"]
-    return dsets, model, global_step
-
-
-if __name__ == "__main__":
-    sys.path.append(os.getcwd())
-
-    parser = get_parser()
-
-    opt, unknown = parser.parse_known_args()
-
-    ckpt = None
-    if opt.resume:
-        if not os.path.exists(opt.resume):
-            raise ValueError("Cannot find {}".format(opt.resume))
-        if os.path.isfile(opt.resume):
-            paths = opt.resume.split("/")
-            try:
-                idx = len(paths)-paths[::-1].index("logs")+1
-            except ValueError:
-                idx = -2 # take a guess: path/to/logdir/checkpoints/model.ckpt
-            logdir = "/".join(paths[:idx])
-            ckpt = opt.resume
-        else:
-            assert os.path.isdir(opt.resume), opt.resume
-            logdir = opt.resume.rstrip("/")
-            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
-        print(f"logdir:{logdir}")
-        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*-project.yaml")))
-        opt.base = base_configs+opt.base
-
-    if opt.config:
-        if type(opt.config) == str:
-            opt.base = [opt.config]
-        else:
-            opt.base = [opt.base[-1]]
-
-    configs = [OmegaConf.load(cfg) for cfg in opt.base]
-    cli = OmegaConf.from_dotlist(unknown)
-    if opt.ignore_base_data:
-        for config in configs:
-            if hasattr(config, "data"): del config["data"]
-    config = OmegaConf.merge(*configs, cli)
-
-    print(ckpt)
-    gpu = True
-    eval_mode = True
-    show_config = False
-    if show_config:
-        print(OmegaConf.to_container(config))
-
-    dsets, model, global_step = load_model_and_dset(config, ckpt, gpu, eval_mode)
-    print(f"Global step: {global_step}")
-
-    outdir = os.path.join(opt.outdir, "{:06}_{}_{}".format(global_step,
-                                                           opt.top_k,
-                                                           opt.temperature))
-    os.makedirs(outdir, exist_ok=True)
-    print("Writing samples to ", outdir)
-    for k in ["originals", "reconstructions", "samples"]:
-        os.makedirs(os.path.join(outdir, k), exist_ok=True)
-    run_conditional(model, dsets, outdir, opt.top_k, opt.temperature)

taming-transformers/scripts/make_scene_samples.py

@@ -1,198 +0,0 @@
-import glob
-import os
-import sys
-from itertools import product
-from pathlib import Path
-from typing import Literal, List, Optional, Tuple
-
-import numpy as np
-import torch
-from omegaconf import OmegaConf
-from pytorch_lightning import seed_everything
-from torch import Tensor
-from torchvision.utils import save_image
-from tqdm import tqdm
-
-from scripts.make_samples import get_parser, load_model_and_dset
-from taming.data.conditional_builder.objects_center_points import ObjectsCenterPointsConditionalBuilder
-from taming.data.helper_types import BoundingBox, Annotation
-from taming.data.annotated_objects_dataset import AnnotatedObjectsDataset
-from taming.models.cond_transformer import Net2NetTransformer
-
-seed_everything(42424242)
-device: Literal['cuda', 'cpu'] = 'cuda'
-first_stage_factor = 16
-trained_on_res = 256
-
-
-def _helper(coord: int, coord_max: int, coord_window: int) -> (int, int):
-    assert 0 <= coord < coord_max
-    coord_desired_center = (coord_window - 1) // 2
-    return np.clip(coord - coord_desired_center, 0, coord_max - coord_window)
-
-
-def get_crop_coordinates(x: int, y: int) -> BoundingBox:
-    WIDTH, HEIGHT = desired_z_shape[1], desired_z_shape[0]
-    x0 = _helper(x, WIDTH, first_stage_factor) / WIDTH
-    y0 = _helper(y, HEIGHT, first_stage_factor) / HEIGHT
-    w = first_stage_factor / WIDTH
-    h = first_stage_factor / HEIGHT
-    return x0, y0, w, h
-
-
-def get_z_indices_crop_out(z_indices: Tensor, predict_x: int, predict_y: int) -> Tensor:
-    WIDTH, HEIGHT = desired_z_shape[1], desired_z_shape[0]
-    x0 = _helper(predict_x, WIDTH, first_stage_factor)
-    y0 = _helper(predict_y, HEIGHT, first_stage_factor)
-    no_images = z_indices.shape[0]
-    cut_out_1 = z_indices[:, y0:predict_y, x0:x0+first_stage_factor].reshape((no_images, -1))
-    cut_out_2 = z_indices[:, predict_y, x0:predict_x]
-    return torch.cat((cut_out_1, cut_out_2), dim=1)
-
-
-@torch.no_grad()
-def sample(model: Net2NetTransformer, annotations: List[Annotation], dataset: AnnotatedObjectsDataset,
-           conditional_builder: ObjectsCenterPointsConditionalBuilder, no_samples: int,
-           temperature: float, top_k: int) -> Tensor:
-    x_max, y_max = desired_z_shape[1], desired_z_shape[0]
-
-    annotations = [a._replace(category_no=dataset.get_category_number(a.category_id)) for a in annotations]
-
-    recompute_conditional = any((desired_resolution[0] > trained_on_res, desired_resolution[1] > trained_on_res))
-    if not recompute_conditional:
-        crop_coordinates = get_crop_coordinates(0, 0)
-        conditional_indices = conditional_builder.build(annotations, crop_coordinates)
-        c_indices = conditional_indices.to(device).repeat(no_samples, 1)
-        z_indices = torch.zeros((no_samples, 0), device=device).long()
-        output_indices = model.sample(z_indices, c_indices, steps=x_max*y_max, temperature=temperature,
-                                      sample=True, top_k=top_k)
-    else:
-        output_indices = torch.zeros((no_samples, y_max, x_max), device=device).long()
-        for predict_y, predict_x in tqdm(product(range(y_max), range(x_max)), desc='sampling_image', total=x_max*y_max):
-            crop_coordinates = get_crop_coordinates(predict_x, predict_y)
-            z_indices = get_z_indices_crop_out(output_indices, predict_x, predict_y)
-            conditional_indices = conditional_builder.build(annotations, crop_coordinates)
-            c_indices = conditional_indices.to(device).repeat(no_samples, 1)
-            new_index = model.sample(z_indices, c_indices, steps=1, temperature=temperature, sample=True, top_k=top_k)
-            output_indices[:, predict_y, predict_x] = new_index[:, -1]
-    z_shape = (
-        no_samples,
-        model.first_stage_model.quantize.e_dim,  # codebook embed_dim
-        desired_z_shape[0],  # z_height
-        desired_z_shape[1]  # z_width
-    )
-    x_sample = model.decode_to_img(output_indices, z_shape) * 0.5 + 0.5
-    x_sample = x_sample.to('cpu')
-
-    plotter = conditional_builder.plot
-    figure_size = (x_sample.shape[2], x_sample.shape[3])
-    scene_graph = conditional_builder.build(annotations, (0., 0., 1., 1.))
-    plot = plotter(scene_graph, dataset.get_textual_label_for_category_no, figure_size)
-    return torch.cat((x_sample, plot.unsqueeze(0)))
-
-
-def get_resolution(resolution_str: str) -> (Tuple[int, int], Tuple[int, int]):
-    if not resolution_str.count(',') == 1:
-        raise ValueError("Give resolution as in 'height,width'")
-    res_h, res_w = resolution_str.split(',')
-    res_h = max(int(res_h), trained_on_res)
-    res_w = max(int(res_w), trained_on_res)
-    z_h = int(round(res_h/first_stage_factor))
-    z_w = int(round(res_w/first_stage_factor))
-    return (z_h, z_w), (z_h*first_stage_factor, z_w*first_stage_factor)
-
-
-def add_arg_to_parser(parser):
-    parser.add_argument(
-        "-R",
-        "--resolution",
-        type=str,
-        default='256,256',
-        help=f"give resolution in multiples of {first_stage_factor}, default is '256,256'",
-    )
-    parser.add_argument(
-        "-C",
-        "--conditional",
-        type=str,
-        default='objects_bbox',
-        help=f"objects_bbox or objects_center_points",
-    )
-    parser.add_argument(
-        "-N",
-        "--n_samples_per_layout",
-        type=int,
-        default=4,
-        help=f"how many samples to generate per layout",
-    )
-    return parser
-
-
-if __name__ == "__main__":
-    sys.path.append(os.getcwd())
-
-    parser = get_parser()
-    parser = add_arg_to_parser(parser)
-
-    opt, unknown = parser.parse_known_args()
-
-    ckpt = None
-    if opt.resume:
-        if not os.path.exists(opt.resume):
-            raise ValueError("Cannot find {}".format(opt.resume))
-        if os.path.isfile(opt.resume):
-            paths = opt.resume.split("/")
-            try:
-                idx = len(paths)-paths[::-1].index("logs")+1
-            except ValueError:
-                idx = -2  # take a guess: path/to/logdir/checkpoints/model.ckpt
-            logdir = "/".join(paths[:idx])
-            ckpt = opt.resume
-        else:
-            assert os.path.isdir(opt.resume), opt.resume
-            logdir = opt.resume.rstrip("/")
-            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
-        print(f"logdir:{logdir}")
-        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*-project.yaml")))
-        opt.base = base_configs+opt.base
-
-    if opt.config:
-        if type(opt.config) == str:
-            opt.base = [opt.config]
-        else:
-            opt.base = [opt.base[-1]]
-
-    configs = [OmegaConf.load(cfg) for cfg in opt.base]
-    cli = OmegaConf.from_dotlist(unknown)
-    if opt.ignore_base_data:
-        for config in configs:
-            if hasattr(config, "data"):
-                del config["data"]
-    config = OmegaConf.merge(*configs, cli)
-    desired_z_shape, desired_resolution = get_resolution(opt.resolution)
-    conditional = opt.conditional
-
-    print(ckpt)
-    gpu = True
-    eval_mode = True
-    show_config = False
-    if show_config:
-        print(OmegaConf.to_container(config))
-
-    dsets, model, global_step = load_model_and_dset(config, ckpt, gpu, eval_mode)
-    print(f"Global step: {global_step}")
-
-    data_loader = dsets.val_dataloader()
-    print(dsets.datasets["validation"].conditional_builders)
-    conditional_builder = dsets.datasets["validation"].conditional_builders[conditional]
-
-    outdir = Path(opt.outdir).joinpath(f"{global_step:06}_{opt.top_k}_{opt.temperature}")
-    outdir.mkdir(exist_ok=True, parents=True)
-    print("Writing samples to ", outdir)
-
-    p_bar_1 = tqdm(enumerate(iter(data_loader)), desc='batch', total=len(data_loader))
-    for batch_no, batch in p_bar_1:
-        save_img: Optional[Tensor] = None
-        for i, annotations in tqdm(enumerate(batch['annotations']), desc='within_batch', total=data_loader.batch_size):
-            imgs = sample(model, annotations, dsets.datasets["validation"], conditional_builder,
-                          opt.n_samples_per_layout, opt.temperature, opt.top_k)
-            save_image(imgs, outdir.joinpath(f'{batch_no:04}_{i:02}.png'), n_row=opt.n_samples_per_layout+1)

taming-transformers/scripts/sample_conditional.py

@@ -1,355 +0,0 @@
-import argparse, os, sys, glob, math, time
-import torch
-import numpy as np
-from omegaconf import OmegaConf
-import streamlit as st
-from streamlit import caching
-from PIL import Image
-from main import instantiate_from_config, DataModuleFromConfig
-from torch.utils.data import DataLoader
-from torch.utils.data.dataloader import default_collate
-
-
-rescale = lambda x: (x + 1.) / 2.
-
-
-def bchw_to_st(x):
-    return rescale(x.detach().cpu().numpy().transpose(0,2,3,1))
-
-def save_img(xstart, fname):
-    I = (xstart.clip(0,1)[0]*255).astype(np.uint8)
-    Image.fromarray(I).save(fname)
-
-
-
-def get_interactive_image(resize=False):
-    image = st.file_uploader("Input", type=["jpg", "JPEG", "png"])
-    if image is not None:
-        image = Image.open(image)
-        if not image.mode == "RGB":
-            image = image.convert("RGB")
-        image = np.array(image).astype(np.uint8)
-        print("upload image shape: {}".format(image.shape))
-        img = Image.fromarray(image)
-        if resize:
-            img = img.resize((256, 256))
-        image = np.array(img)
-        return image
-
-
-def single_image_to_torch(x, permute=True):
-    assert x is not None, "Please provide an image through the upload function"
-    x = np.array(x)
-    x = torch.FloatTensor(x/255.*2. - 1.)[None,...]
-    if permute:
-        x = x.permute(0, 3, 1, 2)
-    return x
-
-
-def pad_to_M(x, M):
-    hp = math.ceil(x.shape[2]/M)*M-x.shape[2]
-    wp = math.ceil(x.shape[3]/M)*M-x.shape[3]
-    x = torch.nn.functional.pad(x, (0,wp,0,hp,0,0,0,0))
-    return x
-
-@torch.no_grad()
-def run_conditional(model, dsets):
-    if len(dsets.datasets) > 1:
-        split = st.sidebar.radio("Split", sorted(dsets.datasets.keys()))
-        dset = dsets.datasets[split]
-    else:
-        dset = next(iter(dsets.datasets.values()))
-    batch_size = 1
-    start_index = st.sidebar.number_input("Example Index (Size: {})".format(len(dset)), value=0,
-                                          min_value=0,
-                                          max_value=len(dset)-batch_size)
-    indices = list(range(start_index, start_index+batch_size))
-
-    example = default_collate([dset[i] for i in indices])
-
-    x = model.get_input("image", example).to(model.device)
-
-    cond_key = model.cond_stage_key
-    c = model.get_input(cond_key, example).to(model.device)
-
-    scale_factor = st.sidebar.slider("Scale Factor", min_value=0.5, max_value=4.0, step=0.25, value=1.00)
-    if scale_factor != 1.0:
-        x = torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode="bicubic")
-        c = torch.nn.functional.interpolate(c, scale_factor=scale_factor, mode="bicubic")
-
-    quant_z, z_indices = model.encode_to_z(x)
-    quant_c, c_indices = model.encode_to_c(c)
-
-    cshape = quant_z.shape
-
-    xrec = model.first_stage_model.decode(quant_z)
-    st.write("image: {}".format(x.shape))
-    st.image(bchw_to_st(x), clamp=True, output_format="PNG")
-    st.write("image reconstruction: {}".format(xrec.shape))
-    st.image(bchw_to_st(xrec), clamp=True, output_format="PNG")
-
-    if cond_key == "segmentation":
-        # get image from segmentation mask
-        num_classes = c.shape[1]
-        c = torch.argmax(c, dim=1, keepdim=True)
-        c = torch.nn.functional.one_hot(c, num_classes=num_classes)
-        c = c.squeeze(1).permute(0, 3, 1, 2).float()
-        c = model.cond_stage_model.to_rgb(c)
-
-    st.write(f"{cond_key}: {tuple(c.shape)}")
-    st.image(bchw_to_st(c), clamp=True, output_format="PNG")
-
-    idx = z_indices
-
-    half_sample = st.sidebar.checkbox("Image Completion", value=False)
-    if half_sample:
-        start = idx.shape[1]//2
-    else:
-        start = 0
-
-    idx[:,start:] = 0
-    idx = idx.reshape(cshape[0],cshape[2],cshape[3])
-    start_i = start//cshape[3]
-    start_j = start %cshape[3]
-
-    if not half_sample and quant_z.shape == quant_c.shape:
-        st.info("Setting idx to c_indices")
-        idx = c_indices.clone().reshape(cshape[0],cshape[2],cshape[3])
-
-    cidx = c_indices
-    cidx = cidx.reshape(quant_c.shape[0],quant_c.shape[2],quant_c.shape[3])
-
-    xstart = model.decode_to_img(idx[:,:cshape[2],:cshape[3]], cshape)
-    st.image(bchw_to_st(xstart), clamp=True, output_format="PNG")
-
-    temperature = st.number_input("Temperature", value=1.0)
-    top_k = st.number_input("Top k", value=100)
-    sample = st.checkbox("Sample", value=True)
-    update_every = st.number_input("Update every", value=75)
-
-    st.text(f"Sampling shape ({cshape[2]},{cshape[3]})")
-
-    animate = st.checkbox("animate")
-    if animate:
-        import imageio
-        outvid = "sampling.mp4"
-        writer = imageio.get_writer(outvid, fps=25)
-    elapsed_t = st.empty()
-    info = st.empty()
-    st.text("Sampled")
-    if st.button("Sample"):
-        output = st.empty()
-        start_t = time.time()
-        for i in range(start_i,cshape[2]-0):
-            if i <= 8:
-                local_i = i
-            elif cshape[2]-i < 8:
-                local_i = 16-(cshape[2]-i)
-            else:
-                local_i = 8
-            for j in range(start_j,cshape[3]-0):
-                if j <= 8:
-                    local_j = j
-                elif cshape[3]-j < 8:
-                    local_j = 16-(cshape[3]-j)
-                else:
-                    local_j = 8
-
-                i_start = i-local_i
-                i_end = i_start+16
-                j_start = j-local_j
-                j_end = j_start+16
-                elapsed_t.text(f"Time: {time.time() - start_t} seconds")
-                info.text(f"Step: ({i},{j}) | Local: ({local_i},{local_j}) | Crop: ({i_start}:{i_end},{j_start}:{j_end})")
-                patch = idx[:,i_start:i_end,j_start:j_end]
-                patch = patch.reshape(patch.shape[0],-1)
-                cpatch = cidx[:, i_start:i_end, j_start:j_end]
-                cpatch = cpatch.reshape(cpatch.shape[0], -1)
-                patch = torch.cat((cpatch, patch), dim=1)
-                logits,_ = model.transformer(patch[:,:-1])
-                logits = logits[:, -256:, :]
-                logits = logits.reshape(cshape[0],16,16,-1)
-                logits = logits[:,local_i,local_j,:]
-
-                logits = logits/temperature
-
-                if top_k is not None:
-                    logits = model.top_k_logits(logits, top_k)
-                # apply softmax to convert to probabilities
-                probs = torch.nn.functional.softmax(logits, dim=-1)
-                # sample from the distribution or take the most likely
-                if sample:
-                    ix = torch.multinomial(probs, num_samples=1)
-                else:
-                    _, ix = torch.topk(probs, k=1, dim=-1)
-                idx[:,i,j] = ix
-
-                if (i*cshape[3]+j)%update_every==0:
-                    xstart = model.decode_to_img(idx[:, :cshape[2], :cshape[3]], cshape,)
-
-                    xstart = bchw_to_st(xstart)
-                    output.image(xstart, clamp=True, output_format="PNG")
-
-                    if animate:
-                        writer.append_data((xstart[0]*255).clip(0, 255).astype(np.uint8))
-
-        xstart = model.decode_to_img(idx[:,:cshape[2],:cshape[3]], cshape)
-        xstart = bchw_to_st(xstart)
-        output.image(xstart, clamp=True, output_format="PNG")
-        #save_img(xstart, "full_res_sample.png")
-        if animate:
-            writer.close()
-            st.video(outvid)
-
-
-def get_parser():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-r",
-        "--resume",
-        type=str,
-        nargs="?",
-        help="load from logdir or checkpoint in logdir",
-    )
-    parser.add_argument(
-        "-b",
-        "--base",
-        nargs="*",
-        metavar="base_config.yaml",
-        help="paths to base configs. Loaded from left-to-right. "
-        "Parameters can be overwritten or added with command-line options of the form `--key value`.",
-        default=list(),
-    )
-    parser.add_argument(
-        "-c",
-        "--config",
-        nargs="?",
-        metavar="single_config.yaml",
-        help="path to single config. If specified, base configs will be ignored "
-        "(except for the last one if left unspecified).",
-        const=True,
-        default="",
-    )
-    parser.add_argument(
-        "--ignore_base_data",
-        action="store_true",
-        help="Ignore data specification from base configs. Useful if you want "
-        "to specify a custom datasets on the command line.",
-    )
-    return parser
-
-
-def load_model_from_config(config, sd, gpu=True, eval_mode=True):
-    if "ckpt_path" in config.params:
-        st.warning("Deleting the restore-ckpt path from the config...")
-        config.params.ckpt_path = None
-    if "downsample_cond_size" in config.params:
-        st.warning("Deleting downsample-cond-size from the config and setting factor=0.5 instead...")
-        config.params.downsample_cond_size = -1
-        config.params["downsample_cond_factor"] = 0.5
-    try:
-        if "ckpt_path" in config.params.first_stage_config.params:
-            config.params.first_stage_config.params.ckpt_path = None
-            st.warning("Deleting the first-stage restore-ckpt path from the config...")
-        if "ckpt_path" in config.params.cond_stage_config.params:
-            config.params.cond_stage_config.params.ckpt_path = None
-            st.warning("Deleting the cond-stage restore-ckpt path from the config...")
-    except:
-        pass
-
-    model = instantiate_from_config(config)
-    if sd is not None:
-        missing, unexpected = model.load_state_dict(sd, strict=False)
-        st.info(f"Missing Keys in State Dict: {missing}")
-        st.info(f"Unexpected Keys in State Dict: {unexpected}")
-    if gpu:
-        model.cuda()
-    if eval_mode:
-        model.eval()
-    return {"model": model}
-
-
-def get_data(config):
-    # get data
-    data = instantiate_from_config(config.data)
-    data.prepare_data()
-    data.setup()
-    return data
-
-
-@st.cache(allow_output_mutation=True, suppress_st_warning=True)
-def load_model_and_dset(config, ckpt, gpu, eval_mode):
-    # get data
-    dsets = get_data(config)   # calls data.config ...
-
-    # now load the specified checkpoint
-    if ckpt:
-        pl_sd = torch.load(ckpt, map_location="cpu")
-        global_step = pl_sd["global_step"]
-    else:
-        pl_sd = {"state_dict": None}
-        global_step = None
-    model = load_model_from_config(config.model,
-                                   pl_sd["state_dict"],
-                                   gpu=gpu,
-                                   eval_mode=eval_mode)["model"]
-    return dsets, model, global_step
-
-
-if __name__ == "__main__":
-    sys.path.append(os.getcwd())
-
-    parser = get_parser()
-
-    opt, unknown = parser.parse_known_args()
-
-    ckpt = None
-    if opt.resume:
-        if not os.path.exists(opt.resume):
-            raise ValueError("Cannot find {}".format(opt.resume))
-        if os.path.isfile(opt.resume):
-            paths = opt.resume.split("/")
-            try:
-                idx = len(paths)-paths[::-1].index("logs")+1
-            except ValueError:
-                idx = -2 # take a guess: path/to/logdir/checkpoints/model.ckpt
-            logdir = "/".join(paths[:idx])
-            ckpt = opt.resume
-        else:
-            assert os.path.isdir(opt.resume), opt.resume
-            logdir = opt.resume.rstrip("/")
-            ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
-        print(f"logdir:{logdir}")
-        base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*-project.yaml")))
-        opt.base = base_configs+opt.base
-
-    if opt.config:
-        if type(opt.config) == str:
-            opt.base = [opt.config]
-        else:
-            opt.base = [opt.base[-1]]
-
-    configs = [OmegaConf.load(cfg) for cfg in opt.base]
-    cli = OmegaConf.from_dotlist(unknown)
-    if opt.ignore_base_data:
-        for config in configs:
-            if hasattr(config, "data"): del config["data"]
-    config = OmegaConf.merge(*configs, cli)
-
-    st.sidebar.text(ckpt)
-    gs = st.sidebar.empty()
-    gs.text(f"Global step: ?")
-    st.sidebar.text("Options")
-    #gpu = st.sidebar.checkbox("GPU", value=True)
-    gpu = True
-    #eval_mode = st.sidebar.checkbox("Eval Mode", value=True)
-    eval_mode = True
-    #show_config = st.sidebar.checkbox("Show Config", value=False)
-    show_config = False
-    if show_config:
-        st.info("Checkpoint: {}".format(ckpt))
-        st.json(OmegaConf.to_container(config))
-
-    dsets, model, global_step = load_model_and_dset(config, ckpt, gpu, eval_mode)
-    gs.text(f"Global step: {global_step}")
-    run_conditional(model, dsets)

taming-transformers/scripts/sample_fast.py

@@ -1,260 +0,0 @@
-import argparse, os, sys, glob
-import torch
-import time
-import numpy as np
-from omegaconf import OmegaConf
-from PIL import Image
-from tqdm import tqdm, trange
-from einops import repeat
-
-from main import instantiate_from_config
-from taming.modules.transformer.mingpt import sample_with_past
-
-
-rescale = lambda x: (x + 1.) / 2.
-
-
-def chw_to_pillow(x):
-    return Image.fromarray((255*rescale(x.detach().cpu().numpy().transpose(1,2,0))).clip(0,255).astype(np.uint8))
-
-
-@torch.no_grad()
-def sample_classconditional(model, batch_size, class_label, steps=256, temperature=None, top_k=None, callback=None,
-                            dim_z=256, h=16, w=16, verbose_time=False, top_p=None):
-    log = dict()
-    assert type(class_label) == int, f'expecting type int but type is {type(class_label)}'
-    qzshape = [batch_size, dim_z, h, w]
-    assert not model.be_unconditional, 'Expecting a class-conditional Net2NetTransformer.'
-    c_indices = repeat(torch.tensor([class_label]), '1 -> b 1', b=batch_size).to(model.device)  # class token
-    t1 = time.time()
-    index_sample = sample_with_past(c_indices, model.transformer, steps=steps,
-                                    sample_logits=True, top_k=top_k, callback=callback,
-                                    temperature=temperature, top_p=top_p)
-    if verbose_time:
-        sampling_time = time.time() - t1
-        print(f"Full sampling takes about {sampling_time:.2f} seconds.")
-    x_sample = model.decode_to_img(index_sample, qzshape)
-    log["samples"] = x_sample
-    log["class_label"] = c_indices
-    return log
-
-
-@torch.no_grad()
-def sample_unconditional(model, batch_size, steps=256, temperature=None, top_k=None, top_p=None, callback=None,
-                         dim_z=256, h=16, w=16, verbose_time=False):
-    log = dict()
-    qzshape = [batch_size, dim_z, h, w]
-    assert model.be_unconditional, 'Expecting an unconditional model.'
-    c_indices = repeat(torch.tensor([model.sos_token]), '1 -> b 1', b=batch_size).to(model.device)  # sos token
-    t1 = time.time()
-    index_sample = sample_with_past(c_indices, model.transformer, steps=steps,
-                                    sample_logits=True, top_k=top_k, callback=callback,
-                                    temperature=temperature, top_p=top_p)
-    if verbose_time:
-        sampling_time = time.time() - t1
-        print(f"Full sampling takes about {sampling_time:.2f} seconds.")
-    x_sample = model.decode_to_img(index_sample, qzshape)
-    log["samples"] = x_sample
-    return log
-
-
-@torch.no_grad()
-def run(logdir, model, batch_size, temperature, top_k, unconditional=True, num_samples=50000,
-        given_classes=None, top_p=None):
-    batches = [batch_size for _ in range(num_samples//batch_size)] + [num_samples % batch_size]
-    if not unconditional:
-        assert given_classes is not None
-        print("Running in pure class-conditional sampling mode. I will produce "
-              f"{num_samples} samples for each of the {len(given_classes)} classes, "
-              f"i.e. {num_samples*len(given_classes)} in total.")
-        for class_label in tqdm(given_classes, desc="Classes"):
-            for n, bs in tqdm(enumerate(batches), desc="Sampling Class"):
-                if bs == 0: break
-                logs = sample_classconditional(model, batch_size=bs, class_label=class_label,
-                                               temperature=temperature, top_k=top_k, top_p=top_p)
-                save_from_logs(logs, logdir, base_count=n * batch_size, cond_key=logs["class_label"])
-    else:
-        print(f"Running in unconditional sampling mode, producing {num_samples} samples.")
-        for n, bs in tqdm(enumerate(batches), desc="Sampling"):
-            if bs == 0: break
-            logs = sample_unconditional(model, batch_size=bs, temperature=temperature, top_k=top_k, top_p=top_p)
-            save_from_logs(logs, logdir, base_count=n * batch_size)
-
-
-def save_from_logs(logs, logdir, base_count, key="samples", cond_key=None):
-    xx = logs[key]
-    for i, x in enumerate(xx):
-        x = chw_to_pillow(x)
-        count = base_count + i
-        if cond_key is None:
-            x.save(os.path.join(logdir, f"{count:06}.png"))
-        else:
-            condlabel = cond_key[i]
-            if type(condlabel) == torch.Tensor: condlabel = condlabel.item()
-            os.makedirs(os.path.join(logdir, str(condlabel)), exist_ok=True)
-            x.save(os.path.join(logdir, str(condlabel), f"{count:06}.png"))
-
-
-def get_parser():
-    def str2bool(v):
-        if isinstance(v, bool):
-            return v
-        if v.lower() in ("yes", "true", "t", "y", "1"):
-            return True
-        elif v.lower() in ("no", "false", "f", "n", "0"):
-            return False
-        else:
-            raise argparse.ArgumentTypeError("Boolean value expected.")
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "-r",
-        "--resume",
-        type=str,
-        nargs="?",
-        help="load from logdir or checkpoint in logdir",
-    )
-    parser.add_argument(
-        "-o",
-        "--outdir",
-        type=str,
-        nargs="?",
-        help="path where the samples will be logged to.",
-        default=""
-    )
-    parser.add_argument(
-        "-b",
-        "--base",
-        nargs="*",
-        metavar="base_config.yaml",
-        help="paths to base configs. Loaded from left-to-right. "
-        "Parameters can be overwritten or added with command-line options of the form `--key value`.",
-        default=list(),
-    )
-    parser.add_argument(
-        "-n",
-        "--num_samples",
-        type=int,
-        nargs="?",
-        help="num_samples to draw",
-        default=50000
-    )
-    parser.add_argument(
-        "--batch_size",
-        type=int,
-        nargs="?",
-        help="the batch size",
-        default=25
-    )
-    parser.add_argument(
-        "-k",
-        "--top_k",
-        type=int,
-        nargs="?",
-        help="top-k value to sample with",
-        default=250,
-    )
-    parser.add_argument(
-        "-t",
-        "--temperature",
-        type=float,
-        nargs="?",
-        help="temperature value to sample with",
-        default=1.0
-    )
-    parser.add_argument(
-        "-p",
-        "--top_p",
-        type=float,
-        nargs="?",
-        help="top-p value to sample with",
-        default=1.0
-    )
-    parser.add_argument(
-        "--classes",
-        type=str,
-        nargs="?",
-        help="specify comma-separated classes to sample from. Uses 1000 classes per default.",
-        default="imagenet"
-    )
-    return parser
-
-
-def load_model_from_config(config, sd, gpu=True, eval_mode=True):
-    model = instantiate_from_config(config)
-    if sd is not None:
-        model.load_state_dict(sd)
-    if gpu:
-        model.cuda()
-    if eval_mode:
-        model.eval()
-    return {"model": model}
-
-
-def load_model(config, ckpt, gpu, eval_mode):
-    # load the specified checkpoint
-    if ckpt:
-        pl_sd = torch.load(ckpt, map_location="cpu")
-        global_step = pl_sd["global_step"]
-        print(f"loaded model from global step {global_step}.")
-    else:
-        pl_sd = {"state_dict": None}
-        global_step = None
-    model = load_model_from_config(config.model, pl_sd["state_dict"], gpu=gpu, eval_mode=eval_mode)["model"]
-    return model, global_step
-
-
-if __name__ == "__main__":
-    sys.path.append(os.getcwd())
-    parser = get_parser()
-
-    opt, unknown = parser.parse_known_args()
-    assert opt.resume
-
-    ckpt = None
-
-    if not os.path.exists(opt.resume):
-        raise ValueError("Cannot find {}".format(opt.resume))
-    if os.path.isfile(opt.resume):
-        paths = opt.resume.split("/")
-        try:
-            idx = len(paths)-paths[::-1].index("logs")+1
-        except ValueError:
-            idx = -2 # take a guess: path/to/logdir/checkpoints/model.ckpt
-        logdir = "/".join(paths[:idx])
-        ckpt = opt.resume
-    else:
-        assert os.path.isdir(opt.resume), opt.resume
-        logdir = opt.resume.rstrip("/")
-        ckpt = os.path.join(logdir, "checkpoints", "last.ckpt")
-
-    base_configs = sorted(glob.glob(os.path.join(logdir, "configs/*-project.yaml")))
-    opt.base = base_configs+opt.base
-
-    configs = [OmegaConf.load(cfg) for cfg in opt.base]
-    cli = OmegaConf.from_dotlist(unknown)
-    config = OmegaConf.merge(*configs, cli)
-
-    model, global_step = load_model(config, ckpt, gpu=True, eval_mode=True)
-
-    if opt.outdir:
-        print(f"Switching logdir from '{logdir}' to '{opt.outdir}'")
-        logdir = opt.outdir
-
-    if opt.classes == "imagenet":
-        given_classes = [i for i in range(1000)]
-    else:
-        cls_str = opt.classes
-        assert not cls_str.endswith(","), 'class string should not end with a ","'
-        given_classes = [int(c) for c in cls_str.split(",")]
-
-    logdir = os.path.join(logdir, "samples", f"top_k_{opt.top_k}_temp_{opt.temperature:.2f}_top_p_{opt.top_p}",
-                          f"{global_step}")
-
-    print(f"Logging to {logdir}")
-    os.makedirs(logdir, exist_ok=True)
-
-    run(logdir, model, opt.batch_size, opt.temperature, opt.top_k, unconditional=model.be_unconditional,
-        given_classes=given_classes, num_samples=opt.num_samples, top_p=opt.top_p)
-
-    print("done.")

taming-transformers/setup.py

@@ -1,13 +0,0 @@
-from setuptools import setup, find_packages
-
-setup(
-    name='taming-transformers',
-    version='0.0.1',
-    description='Taming Transformers for High-Resolution Image Synthesis',
-    packages=find_packages(),
-    install_requires=[
-        'torch',
-        'numpy',
-        'tqdm',
-    ],
-)

taming-transformers/taming/lr_scheduler.py

@@ -1,34 +0,0 @@
-import numpy as np
-
-
-class LambdaWarmUpCosineScheduler:
-    """
-    note: use with a base_lr of 1.0
-    """
-    def __init__(self, warm_up_steps, lr_min, lr_max, lr_start, max_decay_steps, verbosity_interval=0):
-        self.lr_warm_up_steps = warm_up_steps
-        self.lr_start = lr_start
-        self.lr_min = lr_min
-        self.lr_max = lr_max
-        self.lr_max_decay_steps = max_decay_steps
-        self.last_lr = 0.
-        self.verbosity_interval = verbosity_interval
-
-    def schedule(self, n):
-        if self.verbosity_interval > 0:
-            if n % self.verbosity_interval == 0: print(f"current step: {n}, recent lr-multiplier: {self.last_lr}")
-        if n < self.lr_warm_up_steps:
-            lr = (self.lr_max - self.lr_start) / self.lr_warm_up_steps * n + self.lr_start
-            self.last_lr = lr
-            return lr
-        else:
-            t = (n - self.lr_warm_up_steps) / (self.lr_max_decay_steps - self.lr_warm_up_steps)
-            t = min(t, 1.0)
-            lr = self.lr_min + 0.5 * (self.lr_max - self.lr_min) * (
-                    1 + np.cos(t * np.pi))
-            self.last_lr = lr
-            return lr
-
-    def __call__(self, n):
-        return self.schedule(n)
-

taming-transformers/taming/models/cond_transformer.py

@@ -1,352 +0,0 @@
-import os, math
-import torch
-import torch.nn.functional as F
-import pytorch_lightning as pl
-
-from main import instantiate_from_config
-from taming.modules.util import SOSProvider
-
-
-def disabled_train(self, mode=True):
-    """Overwrite model.train with this function to make sure train/eval mode
-    does not change anymore."""
-    return self
-
-
-class Net2NetTransformer(pl.LightningModule):
-    def __init__(self,
-                 transformer_config,
-                 first_stage_config,
-                 cond_stage_config,
-                 permuter_config=None,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 first_stage_key="image",
-                 cond_stage_key="depth",
-                 downsample_cond_size=-1,
-                 pkeep=1.0,
-                 sos_token=0,
-                 unconditional=False,
-                 ):
-        super().__init__()
-        self.be_unconditional = unconditional
-        self.sos_token = sos_token
-        self.first_stage_key = first_stage_key
-        self.cond_stage_key = cond_stage_key
-        self.init_first_stage_from_ckpt(first_stage_config)
-        self.init_cond_stage_from_ckpt(cond_stage_config)
-        if permuter_config is None:
-            permuter_config = {"target": "taming.modules.transformer.permuter.Identity"}
-        self.permuter = instantiate_from_config(config=permuter_config)
-        self.transformer = instantiate_from_config(config=transformer_config)
-
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        self.downsample_cond_size = downsample_cond_size
-        self.pkeep = pkeep
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        for k in sd.keys():
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    self.print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path}")
-
-    def init_first_stage_from_ckpt(self, config):
-        model = instantiate_from_config(config)
-        model = model.eval()
-        model.train = disabled_train
-        self.first_stage_model = model
-
-    def init_cond_stage_from_ckpt(self, config):
-        if config == "__is_first_stage__":
-            print("Using first stage also as cond stage.")
-            self.cond_stage_model = self.first_stage_model
-        elif config == "__is_unconditional__" or self.be_unconditional:
-            print(f"Using no cond stage. Assuming the training is intended to be unconditional. "
-                  f"Prepending {self.sos_token} as a sos token.")
-            self.be_unconditional = True
-            self.cond_stage_key = self.first_stage_key
-            self.cond_stage_model = SOSProvider(self.sos_token)
-        else:
-            model = instantiate_from_config(config)
-            model = model.eval()
-            model.train = disabled_train
-            self.cond_stage_model = model
-
-    def forward(self, x, c):
-        # one step to produce the logits
-        _, z_indices = self.encode_to_z(x)
-        _, c_indices = self.encode_to_c(c)
-
-        if self.training and self.pkeep < 1.0:
-            mask = torch.bernoulli(self.pkeep*torch.ones(z_indices.shape,
-                                                         device=z_indices.device))
-            mask = mask.round().to(dtype=torch.int64)
-            r_indices = torch.randint_like(z_indices, self.transformer.config.vocab_size)
-            a_indices = mask*z_indices+(1-mask)*r_indices
-        else:
-            a_indices = z_indices
-
-        cz_indices = torch.cat((c_indices, a_indices), dim=1)
-
-        # target includes all sequence elements (no need to handle first one
-        # differently because we are conditioning)
-        target = z_indices
-        # make the prediction
-        logits, _ = self.transformer(cz_indices[:, :-1])
-        # cut off conditioning outputs - output i corresponds to p(z_i | z_{<i}, c)
-        logits = logits[:, c_indices.shape[1]-1:]
-
-        return logits, target
-
-    def top_k_logits(self, logits, k):
-        v, ix = torch.topk(logits, k)
-        out = logits.clone()
-        out[out < v[..., [-1]]] = -float('Inf')
-        return out
-
-    @torch.no_grad()
-    def sample(self, x, c, steps, temperature=1.0, sample=False, top_k=None,
-               callback=lambda k: None):
-        x = torch.cat((c,x),dim=1)
-        block_size = self.transformer.get_block_size()
-        assert not self.transformer.training
-        if self.pkeep <= 0.0:
-            # one pass suffices since input is pure noise anyway
-            assert len(x.shape)==2
-            noise_shape = (x.shape[0], steps-1)
-            #noise = torch.randint(self.transformer.config.vocab_size, noise_shape).to(x)
-            noise = c.clone()[:,x.shape[1]-c.shape[1]:-1]
-            x = torch.cat((x,noise),dim=1)
-            logits, _ = self.transformer(x)
-            # take all logits for now and scale by temp
-            logits = logits / temperature
-            # optionally crop probabilities to only the top k options
-            if top_k is not None:
-                logits = self.top_k_logits(logits, top_k)
-            # apply softmax to convert to probabilities
-            probs = F.softmax(logits, dim=-1)
-            # sample from the distribution or take the most likely
-            if sample:
-                shape = probs.shape
-                probs = probs.reshape(shape[0]*shape[1],shape[2])
-                ix = torch.multinomial(probs, num_samples=1)
-                probs = probs.reshape(shape[0],shape[1],shape[2])
-                ix = ix.reshape(shape[0],shape[1])
-            else:
-                _, ix = torch.topk(probs, k=1, dim=-1)
-            # cut off conditioning
-            x = ix[:, c.shape[1]-1:]
-        else:
-            for k in range(steps):
-                callback(k)
-                assert x.size(1) <= block_size # make sure model can see conditioning
-                x_cond = x if x.size(1) <= block_size else x[:, -block_size:]  # crop context if needed
-                logits, _ = self.transformer(x_cond)
-                # pluck the logits at the final step and scale by temperature
-                logits = logits[:, -1, :] / temperature
-                # optionally crop probabilities to only the top k options
-                if top_k is not None:
-                    logits = self.top_k_logits(logits, top_k)
-                # apply softmax to convert to probabilities
-                probs = F.softmax(logits, dim=-1)
-                # sample from the distribution or take the most likely
-                if sample:
-                    ix = torch.multinomial(probs, num_samples=1)
-                else:
-                    _, ix = torch.topk(probs, k=1, dim=-1)
-                # append to the sequence and continue
-                x = torch.cat((x, ix), dim=1)
-            # cut off conditioning
-            x = x[:, c.shape[1]:]
-        return x
-
-    @torch.no_grad()
-    def encode_to_z(self, x):
-        quant_z, _, info = self.first_stage_model.encode(x)
-        indices = info[2].view(quant_z.shape[0], -1)
-        indices = self.permuter(indices)
-        return quant_z, indices
-
-    @torch.no_grad()
-    def encode_to_c(self, c):
-        if self.downsample_cond_size > -1:
-            c = F.interpolate(c, size=(self.downsample_cond_size, self.downsample_cond_size))
-        quant_c, _, [_,_,indices] = self.cond_stage_model.encode(c)
-        if len(indices.shape) > 2:
-            indices = indices.view(c.shape[0], -1)
-        return quant_c, indices
-
-    @torch.no_grad()
-    def decode_to_img(self, index, zshape):
-        index = self.permuter(index, reverse=True)
-        bhwc = (zshape[0],zshape[2],zshape[3],zshape[1])
-        quant_z = self.first_stage_model.quantize.get_codebook_entry(
-            index.reshape(-1), shape=bhwc)
-        x = self.first_stage_model.decode(quant_z)
-        return x
-
-    @torch.no_grad()
-    def log_images(self, batch, temperature=None, top_k=None, callback=None, lr_interface=False, **kwargs):
-        log = dict()
-
-        N = 4
-        if lr_interface:
-            x, c = self.get_xc(batch, N, diffuse=False, upsample_factor=8)
-        else:
-            x, c = self.get_xc(batch, N)
-        x = x.to(device=self.device)
-        c = c.to(device=self.device)
-
-        quant_z, z_indices = self.encode_to_z(x)
-        quant_c, c_indices = self.encode_to_c(c)
-
-        # create a "half"" sample
-        z_start_indices = z_indices[:,:z_indices.shape[1]//2]
-        index_sample = self.sample(z_start_indices, c_indices,
-                                   steps=z_indices.shape[1]-z_start_indices.shape[1],
-                                   temperature=temperature if temperature is not None else 1.0,
-                                   sample=True,
-                                   top_k=top_k if top_k is not None else 100,
-                                   callback=callback if callback is not None else lambda k: None)
-        x_sample = self.decode_to_img(index_sample, quant_z.shape)
-
-        # sample
-        z_start_indices = z_indices[:, :0]
-        index_sample = self.sample(z_start_indices, c_indices,
-                                   steps=z_indices.shape[1],
-                                   temperature=temperature if temperature is not None else 1.0,
-                                   sample=True,
-                                   top_k=top_k if top_k is not None else 100,
-                                   callback=callback if callback is not None else lambda k: None)
-        x_sample_nopix = self.decode_to_img(index_sample, quant_z.shape)
-
-        # det sample
-        z_start_indices = z_indices[:, :0]
-        index_sample = self.sample(z_start_indices, c_indices,
-                                   steps=z_indices.shape[1],
-                                   sample=False,
-                                   callback=callback if callback is not None else lambda k: None)
-        x_sample_det = self.decode_to_img(index_sample, quant_z.shape)
-
-        # reconstruction
-        x_rec = self.decode_to_img(z_indices, quant_z.shape)
-
-        log["inputs"] = x
-        log["reconstructions"] = x_rec
-
-        if self.cond_stage_key in ["objects_bbox", "objects_center_points"]:
-            figure_size = (x_rec.shape[2], x_rec.shape[3])
-            dataset = kwargs["pl_module"].trainer.datamodule.datasets["validation"]
-            label_for_category_no = dataset.get_textual_label_for_category_no
-            plotter = dataset.conditional_builders[self.cond_stage_key].plot
-            log["conditioning"] = torch.zeros_like(log["reconstructions"])
-            for i in range(quant_c.shape[0]):
-                log["conditioning"][i] = plotter(quant_c[i], label_for_category_no, figure_size)
-            log["conditioning_rec"] = log["conditioning"]
-        elif self.cond_stage_key != "image":
-            cond_rec = self.cond_stage_model.decode(quant_c)
-            if self.cond_stage_key == "segmentation":
-                # get image from segmentation mask
-                num_classes = cond_rec.shape[1]
-
-                c = torch.argmax(c, dim=1, keepdim=True)
-                c = F.one_hot(c, num_classes=num_classes)
-                c = c.squeeze(1).permute(0, 3, 1, 2).float()
-                c = self.cond_stage_model.to_rgb(c)
-
-                cond_rec = torch.argmax(cond_rec, dim=1, keepdim=True)
-                cond_rec = F.one_hot(cond_rec, num_classes=num_classes)
-                cond_rec = cond_rec.squeeze(1).permute(0, 3, 1, 2).float()
-                cond_rec = self.cond_stage_model.to_rgb(cond_rec)
-            log["conditioning_rec"] = cond_rec
-            log["conditioning"] = c
-
-        log["samples_half"] = x_sample
-        log["samples_nopix"] = x_sample_nopix
-        log["samples_det"] = x_sample_det
-        return log
-
-    def get_input(self, key, batch):
-        x = batch[key]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        if len(x.shape) == 4:
-            x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
-        if x.dtype == torch.double:
-            x = x.float()
-        return x
-
-    def get_xc(self, batch, N=None):
-        x = self.get_input(self.first_stage_key, batch)
-        c = self.get_input(self.cond_stage_key, batch)
-        if N is not None:
-            x = x[:N]
-            c = c[:N]
-        return x, c
-
-    def shared_step(self, batch, batch_idx):
-        x, c = self.get_xc(batch)
-        logits, target = self(x, c)
-        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), target.reshape(-1))
-        return loss
-
-    def training_step(self, batch, batch_idx):
-        loss = self.shared_step(batch, batch_idx)
-        self.log("train/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        return loss
-
-    def validation_step(self, batch, batch_idx):
-        loss = self.shared_step(batch, batch_idx)
-        self.log("val/loss", loss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        return loss
-
-    def configure_optimizers(self):
-        """
-        Following minGPT:
-        This long function is unfortunately doing something very simple and is being very defensive:
-        We are separating out all parameters of the model into two buckets: those that will experience
-        weight decay for regularization and those that won't (biases, and layernorm/embedding weights).
-        We are then returning the PyTorch optimizer object.
-        """
-        # separate out all parameters to those that will and won't experience regularizing weight decay
-        decay = set()
-        no_decay = set()
-        whitelist_weight_modules = (torch.nn.Linear, )
-        blacklist_weight_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
-        for mn, m in self.transformer.named_modules():
-            for pn, p in m.named_parameters():
-                fpn = '%s.%s' % (mn, pn) if mn else pn # full param name
-
-                if pn.endswith('bias'):
-                    # all biases will not be decayed
-                    no_decay.add(fpn)
-                elif pn.endswith('weight') and isinstance(m, whitelist_weight_modules):
-                    # weights of whitelist modules will be weight decayed
-                    decay.add(fpn)
-                elif pn.endswith('weight') and isinstance(m, blacklist_weight_modules):
-                    # weights of blacklist modules will NOT be weight decayed
-                    no_decay.add(fpn)
-
-        # special case the position embedding parameter in the root GPT module as not decayed
-        no_decay.add('pos_emb')
-
-        # validate that we considered every parameter
-        param_dict = {pn: p for pn, p in self.transformer.named_parameters()}
-        inter_params = decay & no_decay
-        union_params = decay | no_decay
-        assert len(inter_params) == 0, "parameters %s made it into both decay/no_decay sets!" % (str(inter_params), )
-        assert len(param_dict.keys() - union_params) == 0, "parameters %s were not separated into either decay/no_decay set!" \
-                                                    % (str(param_dict.keys() - union_params), )
-
-        # create the pytorch optimizer object
-        optim_groups = [
-            {"params": [param_dict[pn] for pn in sorted(list(decay))], "weight_decay": 0.01},
-            {"params": [param_dict[pn] for pn in sorted(list(no_decay))], "weight_decay": 0.0},
-        ]
-        optimizer = torch.optim.AdamW(optim_groups, lr=self.learning_rate, betas=(0.9, 0.95))
-        return optimizer

taming-transformers/taming/models/dummy_cond_stage.py

@@ -1,22 +0,0 @@
-from torch import Tensor
-
-
-class DummyCondStage:
-    def __init__(self, conditional_key):
-        self.conditional_key = conditional_key
-        self.train = None
-
-    def eval(self):
-        return self
-
-    @staticmethod
-    def encode(c: Tensor):
-        return c, None, (None, None, c)
-
-    @staticmethod
-    def decode(c: Tensor):
-        return c
-
-    @staticmethod
-    def to_rgb(c: Tensor):
-        return c

taming-transformers/taming/models/vqgan.py

@@ -1,404 +0,0 @@
-import torch
-import torch.nn.functional as F
-import pytorch_lightning as pl
-
-from main import instantiate_from_config
-
-from taming.modules.diffusionmodules.model import Encoder, Decoder
-from taming.modules.vqvae.quantize import VectorQuantizer2 as VectorQuantizer
-from taming.modules.vqvae.quantize import GumbelQuantize
-from taming.modules.vqvae.quantize import EMAVectorQuantizer
-
-class VQModel(pl.LightningModule):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 remap=None,
-                 sane_index_shape=False,  # tell vector quantizer to return indices as bhw
-                 ):
-        super().__init__()
-        self.image_key = image_key
-        self.encoder = Encoder(**ddconfig)
-        self.decoder = Decoder(**ddconfig)
-        self.loss = instantiate_from_config(lossconfig)
-        self.quantize = VectorQuantizer(n_embed, embed_dim, beta=0.25,
-                                        remap=remap, sane_index_shape=sane_index_shape)
-        self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1)
-        self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1)
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-        self.image_key = image_key
-        if colorize_nlabels is not None:
-            assert type(colorize_nlabels)==int
-            self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1))
-        if monitor is not None:
-            self.monitor = monitor
-
-    def init_from_ckpt(self, path, ignore_keys=list()):
-        sd = torch.load(path, map_location="cpu")["state_dict"]
-        keys = list(sd.keys())
-        for k in keys:
-            for ik in ignore_keys:
-                if k.startswith(ik):
-                    print("Deleting key {} from state_dict.".format(k))
-                    del sd[k]
-        self.load_state_dict(sd, strict=False)
-        print(f"Restored from {path}")
-
-    def encode(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        quant, emb_loss, info = self.quantize(h)
-        return quant, emb_loss, info
-
-    def decode(self, quant):
-        quant = self.post_quant_conv(quant)
-        dec = self.decoder(quant)
-        return dec
-
-    def decode_code(self, code_b):
-        quant_b = self.quantize.embed_code(code_b)
-        dec = self.decode(quant_b)
-        return dec
-
-    def forward(self, input):
-        quant, diff, _ = self.encode(input)
-        dec = self.decode(quant)
-        return dec, diff
-
-    def get_input(self, batch, k):
-        x = batch[k]
-        if len(x.shape) == 3:
-            x = x[..., None]
-        x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format)
-        return x.float()
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-
-        if optimizer_idx == 0:
-            # autoencode
-            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-
-            self.log("train/aeloss", aeloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # discriminator
-            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log("train/discloss", discloss, prog_bar=True, logger=True, on_step=True, on_epoch=True)
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-
-        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-        rec_loss = log_dict_ae["val/rec_loss"]
-        self.log("val/rec_loss", rec_loss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        self.log("val/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []
-
-    def get_last_layer(self):
-        return self.decoder.conv_out.weight
-
-    def log_images(self, batch, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        return log
-
-    def to_rgb(self, x):
-        assert self.image_key == "segmentation"
-        if not hasattr(self, "colorize"):
-            self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x))
-        x = F.conv2d(x, weight=self.colorize)
-        x = 2.*(x-x.min())/(x.max()-x.min()) - 1.
-        return x
-
-
-class VQSegmentationModel(VQModel):
-    def __init__(self, n_labels, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.register_buffer("colorize", torch.randn(3, n_labels, 1, 1))
-
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        return opt_ae
-
-    def training_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="train")
-        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        return aeloss
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, split="val")
-        self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        total_loss = log_dict_ae["val/total_loss"]
-        self.log("val/total_loss", total_loss,
-                 prog_bar=True, logger=True, on_step=True, on_epoch=True, sync_dist=True)
-        return aeloss
-
-    @torch.no_grad()
-    def log_images(self, batch, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        xrec, _ = self(x)
-        if x.shape[1] > 3:
-            # colorize with random projection
-            assert xrec.shape[1] > 3
-            # convert logits to indices
-            xrec = torch.argmax(xrec, dim=1, keepdim=True)
-            xrec = F.one_hot(xrec, num_classes=x.shape[1])
-            xrec = xrec.squeeze(1).permute(0, 3, 1, 2).float()
-            x = self.to_rgb(x)
-            xrec = self.to_rgb(xrec)
-        log["inputs"] = x
-        log["reconstructions"] = xrec
-        return log
-
-
-class VQNoDiscModel(VQModel):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None
-                 ):
-        super().__init__(ddconfig=ddconfig, lossconfig=lossconfig, n_embed=n_embed, embed_dim=embed_dim,
-                         ckpt_path=ckpt_path, ignore_keys=ignore_keys, image_key=image_key,
-                         colorize_nlabels=colorize_nlabels)
-
-    def training_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        # autoencode
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="train")
-        output = pl.TrainResult(minimize=aeloss)
-        output.log("train/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-        return output
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, self.global_step, split="val")
-        rec_loss = log_dict_ae["val/rec_loss"]
-        output = pl.EvalResult(checkpoint_on=rec_loss)
-        output.log("val/rec_loss", rec_loss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log("val/aeloss", aeloss,
-                   prog_bar=True, logger=True, on_step=True, on_epoch=True)
-        output.log_dict(log_dict_ae)
-
-        return output
-
-    def configure_optimizers(self):
-        optimizer = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quantize.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=self.learning_rate, betas=(0.5, 0.9))
-        return optimizer
-
-
-class GumbelVQ(VQModel):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 temperature_scheduler_config,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 kl_weight=1e-8,
-                 remap=None,
-                 ):
-
-        z_channels = ddconfig["z_channels"]
-        super().__init__(ddconfig,
-                         lossconfig,
-                         n_embed,
-                         embed_dim,
-                         ckpt_path=None,
-                         ignore_keys=ignore_keys,
-                         image_key=image_key,
-                         colorize_nlabels=colorize_nlabels,
-                         monitor=monitor,
-                         )
-
-        self.loss.n_classes = n_embed
-        self.vocab_size = n_embed
-
-        self.quantize = GumbelQuantize(z_channels, embed_dim,
-                                       n_embed=n_embed,
-                                       kl_weight=kl_weight, temp_init=1.0,
-                                       remap=remap)
-
-        self.temperature_scheduler = instantiate_from_config(temperature_scheduler_config)   # annealing of temp
-
-        if ckpt_path is not None:
-            self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys)
-
-    def temperature_scheduling(self):
-        self.quantize.temperature = self.temperature_scheduler(self.global_step)
-
-    def encode_to_prequant(self, x):
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        return h
-
-    def decode_code(self, code_b):
-        raise NotImplementedError
-
-    def training_step(self, batch, batch_idx, optimizer_idx):
-        self.temperature_scheduling()
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x)
-
-        if optimizer_idx == 0:
-            # autoencode
-            aeloss, log_dict_ae = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-
-            self.log_dict(log_dict_ae, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            self.log("temperature", self.quantize.temperature, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return aeloss
-
-        if optimizer_idx == 1:
-            # discriminator
-            discloss, log_dict_disc = self.loss(qloss, x, xrec, optimizer_idx, self.global_step,
-                                            last_layer=self.get_last_layer(), split="train")
-            self.log_dict(log_dict_disc, prog_bar=False, logger=True, on_step=True, on_epoch=True)
-            return discloss
-
-    def validation_step(self, batch, batch_idx):
-        x = self.get_input(batch, self.image_key)
-        xrec, qloss = self(x, return_pred_indices=True)
-        aeloss, log_dict_ae = self.loss(qloss, x, xrec, 0, self.global_step,
-                                        last_layer=self.get_last_layer(), split="val")
-
-        discloss, log_dict_disc = self.loss(qloss, x, xrec, 1, self.global_step,
-                                            last_layer=self.get_last_layer(), split="val")
-        rec_loss = log_dict_ae["val/rec_loss"]
-        self.log("val/rec_loss", rec_loss,
-                 prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        self.log("val/aeloss", aeloss,
-                 prog_bar=True, logger=True, on_step=False, on_epoch=True, sync_dist=True)
-        self.log_dict(log_dict_ae)
-        self.log_dict(log_dict_disc)
-        return self.log_dict
-
-    def log_images(self, batch, **kwargs):
-        log = dict()
-        x = self.get_input(batch, self.image_key)
-        x = x.to(self.device)
-        # encode
-        h = self.encoder(x)
-        h = self.quant_conv(h)
-        quant, _, _ = self.quantize(h)
-        # decode
-        x_rec = self.decode(quant)
-        log["inputs"] = x
-        log["reconstructions"] = x_rec
-        return log
-
-
-class EMAVQ(VQModel):
-    def __init__(self,
-                 ddconfig,
-                 lossconfig,
-                 n_embed,
-                 embed_dim,
-                 ckpt_path=None,
-                 ignore_keys=[],
-                 image_key="image",
-                 colorize_nlabels=None,
-                 monitor=None,
-                 remap=None,
-                 sane_index_shape=False,  # tell vector quantizer to return indices as bhw
-                 ):
-        super().__init__(ddconfig,
-                         lossconfig,
-                         n_embed,
-                         embed_dim,
-                         ckpt_path=None,
-                         ignore_keys=ignore_keys,
-                         image_key=image_key,
-                         colorize_nlabels=colorize_nlabels,
-                         monitor=monitor,
-                         )
-        self.quantize = EMAVectorQuantizer(n_embed=n_embed,
-                                           embedding_dim=embed_dim,
-                                           beta=0.25,
-                                           remap=remap)
-    def configure_optimizers(self):
-        lr = self.learning_rate
-        #Remove self.quantize from parameter list since it is updated via EMA
-        opt_ae = torch.optim.Adam(list(self.encoder.parameters())+
-                                  list(self.decoder.parameters())+
-                                  list(self.quant_conv.parameters())+
-                                  list(self.post_quant_conv.parameters()),
-                                  lr=lr, betas=(0.5, 0.9))
-        opt_disc = torch.optim.Adam(self.loss.discriminator.parameters(),
-                                    lr=lr, betas=(0.5, 0.9))
-        return [opt_ae, opt_disc], []                                           

taming-transformers/taming/modules/diffusionmodules/model.py

@@ -1,776 +0,0 @@
-# pytorch_diffusion + derived encoder decoder
-import math
-import torch
-import torch.nn as nn
-import numpy as np
-
-
-def get_timestep_embedding(timesteps, embedding_dim):
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models:
-    From Fairseq.
-    Build sinusoidal embeddings.
-    This matches the implementation in tensor2tensor, but differs slightly
-    from the description in Section 3.5 of "Attention Is All You Need".
-    """
-    assert len(timesteps.shape) == 1
-
-    half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
-    emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
-    emb = emb.to(device=timesteps.device)
-    emb = timesteps.float()[:, None] * emb[None, :]
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
-    if embedding_dim % 2 == 1:  # zero pad
-        emb = torch.nn.functional.pad(emb, (0,1,0,0))
-    return emb
-
-
-def nonlinearity(x):
-    # swish
-    return x*torch.sigmoid(x)
-
-
-def Normalize(in_channels):
-    return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
-
-
-class Upsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-        if self.with_conv:
-            x = self.conv(x)
-        return x
-
-
-class Downsample(nn.Module):
-    def __init__(self, in_channels, with_conv):
-        super().__init__()
-        self.with_conv = with_conv
-        if self.with_conv:
-            # no asymmetric padding in torch conv, must do it ourselves
-            self.conv = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=3,
-                                        stride=2,
-                                        padding=0)
-
-    def forward(self, x):
-        if self.with_conv:
-            pad = (0,1,0,1)
-            x = torch.nn.functional.pad(x, pad, mode="constant", value=0)
-            x = self.conv(x)
-        else:
-            x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2)
-        return x
-
-
-class ResnetBlock(nn.Module):
-    def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout, temb_channels=512):
-        super().__init__()
-        self.in_channels = in_channels
-        out_channels = in_channels if out_channels is None else out_channels
-        self.out_channels = out_channels
-        self.use_conv_shortcut = conv_shortcut
-
-        self.norm1 = Normalize(in_channels)
-        self.conv1 = torch.nn.Conv2d(in_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if temb_channels > 0:
-            self.temb_proj = torch.nn.Linear(temb_channels,
-                                             out_channels)
-        self.norm2 = Normalize(out_channels)
-        self.dropout = torch.nn.Dropout(dropout)
-        self.conv2 = torch.nn.Conv2d(out_channels,
-                                     out_channels,
-                                     kernel_size=3,
-                                     stride=1,
-                                     padding=1)
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                self.conv_shortcut = torch.nn.Conv2d(in_channels,
-                                                     out_channels,
-                                                     kernel_size=3,
-                                                     stride=1,
-                                                     padding=1)
-            else:
-                self.nin_shortcut = torch.nn.Conv2d(in_channels,
-                                                    out_channels,
-                                                    kernel_size=1,
-                                                    stride=1,
-                                                    padding=0)
-
-    def forward(self, x, temb):
-        h = x
-        h = self.norm1(h)
-        h = nonlinearity(h)
-        h = self.conv1(h)
-
-        if temb is not None:
-            h = h + self.temb_proj(nonlinearity(temb))[:,:,None,None]
-
-        h = self.norm2(h)
-        h = nonlinearity(h)
-        h = self.dropout(h)
-        h = self.conv2(h)
-
-        if self.in_channels != self.out_channels:
-            if self.use_conv_shortcut:
-                x = self.conv_shortcut(x)
-            else:
-                x = self.nin_shortcut(x)
-
-        return x+h
-
-
-class AttnBlock(nn.Module):
-    def __init__(self, in_channels):
-        super().__init__()
-        self.in_channels = in_channels
-
-        self.norm = Normalize(in_channels)
-        self.q = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.k = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.v = torch.nn.Conv2d(in_channels,
-                                 in_channels,
-                                 kernel_size=1,
-                                 stride=1,
-                                 padding=0)
-        self.proj_out = torch.nn.Conv2d(in_channels,
-                                        in_channels,
-                                        kernel_size=1,
-                                        stride=1,
-                                        padding=0)
-
-
-    def forward(self, x):
-        h_ = x
-        h_ = self.norm(h_)
-        q = self.q(h_)
-        k = self.k(h_)
-        v = self.v(h_)
-
-        # compute attention
-        b,c,h,w = q.shape
-        q = q.reshape(b,c,h*w)
-        q = q.permute(0,2,1)   # b,hw,c
-        k = k.reshape(b,c,h*w) # b,c,hw
-        w_ = torch.bmm(q,k)     # b,hw,hw    w[b,i,j]=sum_c q[b,i,c]k[b,c,j]
-        w_ = w_ * (int(c)**(-0.5))
-        w_ = torch.nn.functional.softmax(w_, dim=2)
-
-        # attend to values
-        v = v.reshape(b,c,h*w)
-        w_ = w_.permute(0,2,1)   # b,hw,hw (first hw of k, second of q)
-        h_ = torch.bmm(v,w_)     # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j]
-        h_ = h_.reshape(b,c,h,w)
-
-        h_ = self.proj_out(h_)
-
-        return x+h_
-
-
-class Model(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, use_timestep=True):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                torch.nn.Linear(self.ch,
-                                self.temb_ch),
-                torch.nn.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x, t=None):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Encoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, double_z=True, **ignore_kwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(in_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        2*z_channels if double_z else z_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x):
-        #assert x.shape[2] == x.shape[3] == self.resolution, "{}, {}, {}".format(x.shape[2], x.shape[3], self.resolution)
-
-        # timestep embedding
-        temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class Decoder(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True, in_channels,
-                 resolution, z_channels, give_pre_end=False, **ignorekwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.give_pre_end = give_pre_end
-
-        # compute in_ch_mult, block_in and curr_res at lowest res
-        in_ch_mult = (1,)+tuple(ch_mult)
-        block_in = ch*ch_mult[self.num_resolutions-1]
-        curr_res = resolution // 2**(self.num_resolutions-1)
-        self.z_shape = (1,z_channels,curr_res,curr_res)
-        print("Working with z of shape {} = {} dimensions.".format(
-            self.z_shape, np.prod(self.z_shape)))
-
-        # z to block_in
-        self.conv_in = torch.nn.Conv2d(z_channels,
-                                       block_in,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, z):
-        #assert z.shape[1:] == self.z_shape[1:]
-        self.last_z_shape = z.shape
-
-        # timestep embedding
-        temb = None
-
-        # z to block_in
-        h = self.conv_in(z)
-
-        # middle
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](h, temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        if self.give_pre_end:
-            return h
-
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class VUNet(nn.Module):
-    def __init__(self, *, ch, out_ch, ch_mult=(1,2,4,8), num_res_blocks,
-                 attn_resolutions, dropout=0.0, resamp_with_conv=True,
-                 in_channels, c_channels,
-                 resolution, z_channels, use_timestep=False, **ignore_kwargs):
-        super().__init__()
-        self.ch = ch
-        self.temb_ch = self.ch*4
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        self.resolution = resolution
-
-        self.use_timestep = use_timestep
-        if self.use_timestep:
-            # timestep embedding
-            self.temb = nn.Module()
-            self.temb.dense = nn.ModuleList([
-                torch.nn.Linear(self.ch,
-                                self.temb_ch),
-                torch.nn.Linear(self.temb_ch,
-                                self.temb_ch),
-            ])
-
-        # downsampling
-        self.conv_in = torch.nn.Conv2d(c_channels,
-                                       self.ch,
-                                       kernel_size=3,
-                                       stride=1,
-                                       padding=1)
-
-        curr_res = resolution
-        in_ch_mult = (1,)+tuple(ch_mult)
-        self.down = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_in = ch*in_ch_mult[i_level]
-            block_out = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks):
-                block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            down = nn.Module()
-            down.block = block
-            down.attn = attn
-            if i_level != self.num_resolutions-1:
-                down.downsample = Downsample(block_in, resamp_with_conv)
-                curr_res = curr_res // 2
-            self.down.append(down)
-
-        self.z_in = torch.nn.Conv2d(z_channels,
-                                    block_in,
-                                    kernel_size=1,
-                                    stride=1,
-                                    padding=0)
-        # middle
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=2*block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-        self.mid.attn_1 = AttnBlock(block_in)
-        self.mid.block_2 = ResnetBlock(in_channels=block_in,
-                                       out_channels=block_in,
-                                       temb_channels=self.temb_ch,
-                                       dropout=dropout)
-
-        # upsampling
-        self.up = nn.ModuleList()
-        for i_level in reversed(range(self.num_resolutions)):
-            block = nn.ModuleList()
-            attn = nn.ModuleList()
-            block_out = ch*ch_mult[i_level]
-            skip_in = ch*ch_mult[i_level]
-            for i_block in range(self.num_res_blocks+1):
-                if i_block == self.num_res_blocks:
-                    skip_in = ch*in_ch_mult[i_level]
-                block.append(ResnetBlock(in_channels=block_in+skip_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-                if curr_res in attn_resolutions:
-                    attn.append(AttnBlock(block_in))
-            up = nn.Module()
-            up.block = block
-            up.attn = attn
-            if i_level != 0:
-                up.upsample = Upsample(block_in, resamp_with_conv)
-                curr_res = curr_res * 2
-            self.up.insert(0, up) # prepend to get consistent order
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_ch,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-
-    def forward(self, x, z):
-        #assert x.shape[2] == x.shape[3] == self.resolution
-
-        if self.use_timestep:
-            # timestep embedding
-            assert t is not None
-            temb = get_timestep_embedding(t, self.ch)
-            temb = self.temb.dense[0](temb)
-            temb = nonlinearity(temb)
-            temb = self.temb.dense[1](temb)
-        else:
-            temb = None
-
-        # downsampling
-        hs = [self.conv_in(x)]
-        for i_level in range(self.num_resolutions):
-            for i_block in range(self.num_res_blocks):
-                h = self.down[i_level].block[i_block](hs[-1], temb)
-                if len(self.down[i_level].attn) > 0:
-                    h = self.down[i_level].attn[i_block](h)
-                hs.append(h)
-            if i_level != self.num_resolutions-1:
-                hs.append(self.down[i_level].downsample(hs[-1]))
-
-        # middle
-        h = hs[-1]
-        z = self.z_in(z)
-        h = torch.cat((h,z),dim=1)
-        h = self.mid.block_1(h, temb)
-        h = self.mid.attn_1(h)
-        h = self.mid.block_2(h, temb)
-
-        # upsampling
-        for i_level in reversed(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks+1):
-                h = self.up[i_level].block[i_block](
-                    torch.cat([h, hs.pop()], dim=1), temb)
-                if len(self.up[i_level].attn) > 0:
-                    h = self.up[i_level].attn[i_block](h)
-            if i_level != 0:
-                h = self.up[i_level].upsample(h)
-
-        # end
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-
-
-class SimpleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, *args, **kwargs):
-        super().__init__()
-        self.model = nn.ModuleList([nn.Conv2d(in_channels, in_channels, 1),
-                                     ResnetBlock(in_channels=in_channels,
-                                                 out_channels=2 * in_channels,
-                                                 temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=2 * in_channels,
-                                                out_channels=4 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     ResnetBlock(in_channels=4 * in_channels,
-                                                out_channels=2 * in_channels,
-                                                temb_channels=0, dropout=0.0),
-                                     nn.Conv2d(2*in_channels, in_channels, 1),
-                                     Upsample(in_channels, with_conv=True)])
-        # end
-        self.norm_out = Normalize(in_channels)
-        self.conv_out = torch.nn.Conv2d(in_channels,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        for i, layer in enumerate(self.model):
-            if i in [1,2,3]:
-                x = layer(x, None)
-            else:
-                x = layer(x)
-
-        h = self.norm_out(x)
-        h = nonlinearity(h)
-        x = self.conv_out(h)
-        return x
-
-
-class UpsampleDecoder(nn.Module):
-    def __init__(self, in_channels, out_channels, ch, num_res_blocks, resolution,
-                 ch_mult=(2,2), dropout=0.0):
-        super().__init__()
-        # upsampling
-        self.temb_ch = 0
-        self.num_resolutions = len(ch_mult)
-        self.num_res_blocks = num_res_blocks
-        block_in = in_channels
-        curr_res = resolution // 2 ** (self.num_resolutions - 1)
-        self.res_blocks = nn.ModuleList()
-        self.upsample_blocks = nn.ModuleList()
-        for i_level in range(self.num_resolutions):
-            res_block = []
-            block_out = ch * ch_mult[i_level]
-            for i_block in range(self.num_res_blocks + 1):
-                res_block.append(ResnetBlock(in_channels=block_in,
-                                         out_channels=block_out,
-                                         temb_channels=self.temb_ch,
-                                         dropout=dropout))
-                block_in = block_out
-            self.res_blocks.append(nn.ModuleList(res_block))
-            if i_level != self.num_resolutions - 1:
-                self.upsample_blocks.append(Upsample(block_in, True))
-                curr_res = curr_res * 2
-
-        # end
-        self.norm_out = Normalize(block_in)
-        self.conv_out = torch.nn.Conv2d(block_in,
-                                        out_channels,
-                                        kernel_size=3,
-                                        stride=1,
-                                        padding=1)
-
-    def forward(self, x):
-        # upsampling
-        h = x
-        for k, i_level in enumerate(range(self.num_resolutions)):
-            for i_block in range(self.num_res_blocks + 1):
-                h = self.res_blocks[i_level][i_block](h, None)
-            if i_level != self.num_resolutions - 1:
-                h = self.upsample_blocks[k](h)
-        h = self.norm_out(h)
-        h = nonlinearity(h)
-        h = self.conv_out(h)
-        return h
-

taming-transformers/taming/modules/discriminator/model.py

@@ -1,67 +0,0 @@
-import functools
-import torch.nn as nn
-
-
-from taming.modules.util import ActNorm
-
-
-def weights_init(m):
-    classname = m.__class__.__name__
-    if classname.find('Conv') != -1:
-        nn.init.normal_(m.weight.data, 0.0, 0.02)
-    elif classname.find('BatchNorm') != -1:
-        nn.init.normal_(m.weight.data, 1.0, 0.02)
-        nn.init.constant_(m.bias.data, 0)
-
-
-class NLayerDiscriminator(nn.Module):
-    """Defines a PatchGAN discriminator as in Pix2Pix
-        --> see https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/models/networks.py
-    """
-    def __init__(self, input_nc=3, ndf=64, n_layers=3, use_actnorm=False):
-        """Construct a PatchGAN discriminator
-        Parameters:
-            input_nc (int)  -- the number of channels in input images
-            ndf (int)       -- the number of filters in the last conv layer
-            n_layers (int)  -- the number of conv layers in the discriminator
-            norm_layer      -- normalization layer
-        """
-        super(NLayerDiscriminator, self).__init__()
-        if not use_actnorm:
-            norm_layer = nn.BatchNorm2d
-        else:
-            norm_layer = ActNorm
-        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters
-            use_bias = norm_layer.func != nn.BatchNorm2d
-        else:
-            use_bias = norm_layer != nn.BatchNorm2d
-
-        kw = 4
-        padw = 1
-        sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]
-        nf_mult = 1
-        nf_mult_prev = 1
-        for n in range(1, n_layers):  # gradually increase the number of filters
-            nf_mult_prev = nf_mult
-            nf_mult = min(2 ** n, 8)
-            sequence += [
-                nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),
-                norm_layer(ndf * nf_mult),
-                nn.LeakyReLU(0.2, True)
-            ]
-
-        nf_mult_prev = nf_mult
-        nf_mult = min(2 ** n_layers, 8)
-        sequence += [
-            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),
-            norm_layer(ndf * nf_mult),
-            nn.LeakyReLU(0.2, True)
-        ]
-
-        sequence += [
-            nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map
-        self.main = nn.Sequential(*sequence)
-
-    def forward(self, input):
-        """Standard forward."""
-        return self.main(input)

taming-transformers/taming/modules/losses/__init__.py

@@ -1,2 +0,0 @@
-from taming.modules.losses.vqperceptual import DummyLoss
-

taming-transformers/taming/modules/losses/lpips.py

@@ -1,123 +0,0 @@
-"""Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models"""
-
-import torch
-import torch.nn as nn
-from torchvision import models
-from collections import namedtuple
-
-from taming.util import get_ckpt_path
-
-
-class LPIPS(nn.Module):
-    # Learned perceptual metric
-    def __init__(self, use_dropout=True):
-        super().__init__()
-        self.scaling_layer = ScalingLayer()
-        self.chns = [64, 128, 256, 512, 512]  # vg16 features
-        self.net = vgg16(pretrained=True, requires_grad=False)
-        self.lin0 = NetLinLayer(self.chns[0], use_dropout=use_dropout)
-        self.lin1 = NetLinLayer(self.chns[1], use_dropout=use_dropout)
-        self.lin2 = NetLinLayer(self.chns[2], use_dropout=use_dropout)
-        self.lin3 = NetLinLayer(self.chns[3], use_dropout=use_dropout)
-        self.lin4 = NetLinLayer(self.chns[4], use_dropout=use_dropout)
-        self.load_from_pretrained()
-        for param in self.parameters():
-            param.requires_grad = False
-
-    def load_from_pretrained(self, name="vgg_lpips"):
-        ckpt = get_ckpt_path(name, "taming/modules/autoencoder/lpips")
-        self.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
-        print("loaded pretrained LPIPS loss from {}".format(ckpt))
-
-    @classmethod
-    def from_pretrained(cls, name="vgg_lpips"):
-        if name != "vgg_lpips":
-            raise NotImplementedError
-        model = cls()
-        ckpt = get_ckpt_path(name)
-        model.load_state_dict(torch.load(ckpt, map_location=torch.device("cpu")), strict=False)
-        return model
-
-    def forward(self, input, target):
-        in0_input, in1_input = (self.scaling_layer(input), self.scaling_layer(target))
-        outs0, outs1 = self.net(in0_input), self.net(in1_input)
-        feats0, feats1, diffs = {}, {}, {}
-        lins = [self.lin0, self.lin1, self.lin2, self.lin3, self.lin4]
-        for kk in range(len(self.chns)):
-            feats0[kk], feats1[kk] = normalize_tensor(outs0[kk]), normalize_tensor(outs1[kk])
-            diffs[kk] = (feats0[kk] - feats1[kk]) ** 2
-
-        res = [spatial_average(lins[kk].model(diffs[kk]), keepdim=True) for kk in range(len(self.chns))]
-        val = res[0]
-        for l in range(1, len(self.chns)):
-            val += res[l]
-        return val
-
-
-class ScalingLayer(nn.Module):
-    def __init__(self):
-        super(ScalingLayer, self).__init__()
-        self.register_buffer('shift', torch.Tensor([-.030, -.088, -.188])[None, :, None, None])
-        self.register_buffer('scale', torch.Tensor([.458, .448, .450])[None, :, None, None])
-
-    def forward(self, inp):
-        return (inp - self.shift) / self.scale
-
-
-class NetLinLayer(nn.Module):
-    """ A single linear layer which does a 1x1 conv """
-    def __init__(self, chn_in, chn_out=1, use_dropout=False):
-        super(NetLinLayer, self).__init__()
-        layers = [nn.Dropout(), ] if (use_dropout) else []
-        layers += [nn.Conv2d(chn_in, chn_out, 1, stride=1, padding=0, bias=False), ]
-        self.model = nn.Sequential(*layers)
-
-
-class vgg16(torch.nn.Module):
-    def __init__(self, requires_grad=False, pretrained=True):
-        super(vgg16, self).__init__()
-        vgg_pretrained_features = models.vgg16(pretrained=pretrained).features
-        self.slice1 = torch.nn.Sequential()
-        self.slice2 = torch.nn.Sequential()
-        self.slice3 = torch.nn.Sequential()
-        self.slice4 = torch.nn.Sequential()
-        self.slice5 = torch.nn.Sequential()
-        self.N_slices = 5
-        for x in range(4):
-            self.slice1.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(4, 9):
-            self.slice2.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(9, 16):
-            self.slice3.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(16, 23):
-            self.slice4.add_module(str(x), vgg_pretrained_features[x])
-        for x in range(23, 30):
-            self.slice5.add_module(str(x), vgg_pretrained_features[x])
-        if not requires_grad:
-            for param in self.parameters():
-                param.requires_grad = False
-
-    def forward(self, X):
-        h = self.slice1(X)
-        h_relu1_2 = h
-        h = self.slice2(h)
-        h_relu2_2 = h
-        h = self.slice3(h)
-        h_relu3_3 = h
-        h = self.slice4(h)
-        h_relu4_3 = h
-        h = self.slice5(h)
-        h_relu5_3 = h
-        vgg_outputs = namedtuple("VggOutputs", ['relu1_2', 'relu2_2', 'relu3_3', 'relu4_3', 'relu5_3'])
-        out = vgg_outputs(h_relu1_2, h_relu2_2, h_relu3_3, h_relu4_3, h_relu5_3)
-        return out
-
-
-def normalize_tensor(x,eps=1e-10):
-    norm_factor = torch.sqrt(torch.sum(x**2,dim=1,keepdim=True))
-    return x/(norm_factor+eps)
-
-
-def spatial_average(x, keepdim=True):
-    return x.mean([2,3],keepdim=keepdim)
-

taming-transformers/taming/modules/losses/segmentation.py

@@ -1,22 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-
-
-class BCELoss(nn.Module):
-    def forward(self, prediction, target):
-        loss = F.binary_cross_entropy_with_logits(prediction,target)
-        return loss, {}
-
-
-class BCELossWithQuant(nn.Module):
-    def __init__(self, codebook_weight=1.):
-        super().__init__()
-        self.codebook_weight = codebook_weight
-
-    def forward(self, qloss, target, prediction, split):
-        bce_loss = F.binary_cross_entropy_with_logits(prediction,target)
-        loss = bce_loss + self.codebook_weight*qloss
-        return loss, {"{}/total_loss".format(split): loss.clone().detach().mean(),
-                      "{}/bce_loss".format(split): bce_loss.detach().mean(),
-                      "{}/quant_loss".format(split): qloss.detach().mean()
-                      }

taming-transformers/taming/modules/losses/vqperceptual.py

@@ -1,136 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from taming.modules.losses.lpips import LPIPS
-from taming.modules.discriminator.model import NLayerDiscriminator, weights_init
-
-
-class DummyLoss(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-
-def adopt_weight(weight, global_step, threshold=0, value=0.):
-    if global_step < threshold:
-        weight = value
-    return weight
-
-
-def hinge_d_loss(logits_real, logits_fake):
-    loss_real = torch.mean(F.relu(1. - logits_real))
-    loss_fake = torch.mean(F.relu(1. + logits_fake))
-    d_loss = 0.5 * (loss_real + loss_fake)
-    return d_loss
-
-
-def vanilla_d_loss(logits_real, logits_fake):
-    d_loss = 0.5 * (
-        torch.mean(torch.nn.functional.softplus(-logits_real)) +
-        torch.mean(torch.nn.functional.softplus(logits_fake)))
-    return d_loss
-
-
-class VQLPIPSWithDiscriminator(nn.Module):
-    def __init__(self, disc_start, codebook_weight=1.0, pixelloss_weight=1.0,
-                 disc_num_layers=3, disc_in_channels=3, disc_factor=1.0, disc_weight=1.0,
-                 perceptual_weight=1.0, use_actnorm=False, disc_conditional=False,
-                 disc_ndf=64, disc_loss="hinge"):
-        super().__init__()
-        assert disc_loss in ["hinge", "vanilla"]
-        self.codebook_weight = codebook_weight
-        self.pixel_weight = pixelloss_weight
-        self.perceptual_loss = LPIPS().eval()
-        self.perceptual_weight = perceptual_weight
-
-        self.discriminator = NLayerDiscriminator(input_nc=disc_in_channels,
-                                                 n_layers=disc_num_layers,
-                                                 use_actnorm=use_actnorm,
-                                                 ndf=disc_ndf
-                                                 ).apply(weights_init)
-        self.discriminator_iter_start = disc_start
-        if disc_loss == "hinge":
-            self.disc_loss = hinge_d_loss
-        elif disc_loss == "vanilla":
-            self.disc_loss = vanilla_d_loss
-        else:
-            raise ValueError(f"Unknown GAN loss '{disc_loss}'.")
-        print(f"VQLPIPSWithDiscriminator running with {disc_loss} loss.")
-        self.disc_factor = disc_factor
-        self.discriminator_weight = disc_weight
-        self.disc_conditional = disc_conditional
-
-    def calculate_adaptive_weight(self, nll_loss, g_loss, last_layer=None):
-        if last_layer is not None:
-            nll_grads = torch.autograd.grad(nll_loss, last_layer, retain_graph=True)[0]
-            g_grads = torch.autograd.grad(g_loss, last_layer, retain_graph=True)[0]
-        else:
-            nll_grads = torch.autograd.grad(nll_loss, self.last_layer[0], retain_graph=True)[0]
-            g_grads = torch.autograd.grad(g_loss, self.last_layer[0], retain_graph=True)[0]
-
-        d_weight = torch.norm(nll_grads) / (torch.norm(g_grads) + 1e-4)
-        d_weight = torch.clamp(d_weight, 0.0, 1e4).detach()
-        d_weight = d_weight * self.discriminator_weight
-        return d_weight
-
-    def forward(self, codebook_loss, inputs, reconstructions, optimizer_idx,
-                global_step, last_layer=None, cond=None, split="train"):
-        rec_loss = torch.abs(inputs.contiguous() - reconstructions.contiguous())
-        if self.perceptual_weight > 0:
-            p_loss = self.perceptual_loss(inputs.contiguous(), reconstructions.contiguous())
-            rec_loss = rec_loss + self.perceptual_weight * p_loss
-        else:
-            p_loss = torch.tensor([0.0])
-
-        nll_loss = rec_loss
-        #nll_loss = torch.sum(nll_loss) / nll_loss.shape[0]
-        nll_loss = torch.mean(nll_loss)
-
-        # now the GAN part
-        if optimizer_idx == 0:
-            # generator update
-            if cond is None:
-                assert not self.disc_conditional
-                logits_fake = self.discriminator(reconstructions.contiguous())
-            else:
-                assert self.disc_conditional
-                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous(), cond), dim=1))
-            g_loss = -torch.mean(logits_fake)
-
-            try:
-                d_weight = self.calculate_adaptive_weight(nll_loss, g_loss, last_layer=last_layer)
-            except RuntimeError:
-                assert not self.training
-                d_weight = torch.tensor(0.0)
-
-            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
-            loss = nll_loss + d_weight * disc_factor * g_loss + self.codebook_weight * codebook_loss.mean()
-
-            log = {"{}/total_loss".format(split): loss.clone().detach().mean(),
-                   "{}/quant_loss".format(split): codebook_loss.detach().mean(),
-                   "{}/nll_loss".format(split): nll_loss.detach().mean(),
-                   "{}/rec_loss".format(split): rec_loss.detach().mean(),
-                   "{}/p_loss".format(split): p_loss.detach().mean(),
-                   "{}/d_weight".format(split): d_weight.detach(),
-                   "{}/disc_factor".format(split): torch.tensor(disc_factor),
-                   "{}/g_loss".format(split): g_loss.detach().mean(),
-                   }
-            return loss, log
-
-        if optimizer_idx == 1:
-            # second pass for discriminator update
-            if cond is None:
-                logits_real = self.discriminator(inputs.contiguous().detach())
-                logits_fake = self.discriminator(reconstructions.contiguous().detach())
-            else:
-                logits_real = self.discriminator(torch.cat((inputs.contiguous().detach(), cond), dim=1))
-                logits_fake = self.discriminator(torch.cat((reconstructions.contiguous().detach(), cond), dim=1))
-
-            disc_factor = adopt_weight(self.disc_factor, global_step, threshold=self.discriminator_iter_start)
-            d_loss = disc_factor * self.disc_loss(logits_real, logits_fake)
-
-            log = {"{}/disc_loss".format(split): d_loss.clone().detach().mean(),
-                   "{}/logits_real".format(split): logits_real.detach().mean(),
-                   "{}/logits_fake".format(split): logits_fake.detach().mean()
-                   }
-            return d_loss, log

taming-transformers/taming/modules/misc/coord.py

@@ -1,31 +0,0 @@
-import torch
-
-class CoordStage(object):
-    def __init__(self, n_embed, down_factor):
-        self.n_embed = n_embed
-        self.down_factor = down_factor
-
-    def eval(self):
-        return self
-
-    def encode(self, c):
-        """fake vqmodel interface"""
-        assert 0.0 <= c.min() and c.max() <= 1.0
-        b,ch,h,w = c.shape
-        assert ch == 1
-
-        c = torch.nn.functional.interpolate(c, scale_factor=1/self.down_factor,
-                                            mode="area")
-        c = c.clamp(0.0, 1.0)
-        c = self.n_embed*c
-        c_quant = c.round()
-        c_ind = c_quant.to(dtype=torch.long)
-
-        info = None, None, c_ind
-        return c_quant, None, info
-
-    def decode(self, c):
-        c = c/self.n_embed
-        c = torch.nn.functional.interpolate(c, scale_factor=self.down_factor,
-                                            mode="nearest")
-        return c

taming-transformers/taming/modules/transformer/mingpt.py

@@ -1,415 +0,0 @@
-"""
-taken from: https://github.com/karpathy/minGPT/
-GPT model:
-- the initial stem consists of a combination of token encoding and a positional encoding
-- the meat of it is a uniform sequence of Transformer blocks
-    - each Transformer is a sequential combination of a 1-hidden-layer MLP block and a self-attention block
-    - all blocks feed into a central residual pathway similar to resnets
-- the final decoder is a linear projection into a vanilla Softmax classifier
-"""
-
-import math
-import logging
-
-import torch
-import torch.nn as nn
-from torch.nn import functional as F
-from transformers import top_k_top_p_filtering
-
-logger = logging.getLogger(__name__)
-
-
-class GPTConfig:
-    """ base GPT config, params common to all GPT versions """
-    embd_pdrop = 0.1
-    resid_pdrop = 0.1
-    attn_pdrop = 0.1
-
-    def __init__(self, vocab_size, block_size, **kwargs):
-        self.vocab_size = vocab_size
-        self.block_size = block_size
-        for k,v in kwargs.items():
-            setattr(self, k, v)
-
-
-class GPT1Config(GPTConfig):
-    """ GPT-1 like network roughly 125M params """
-    n_layer = 12
-    n_head = 12
-    n_embd = 768
-
-
-class CausalSelfAttention(nn.Module):
-    """
-    A vanilla multi-head masked self-attention layer with a projection at the end.
-    It is possible to use torch.nn.MultiheadAttention here but I am including an
-    explicit implementation here to show that there is nothing too scary here.
-    """
-
-    def __init__(self, config):
-        super().__init__()
-        assert config.n_embd % config.n_head == 0
-        # key, query, value projections for all heads
-        self.key = nn.Linear(config.n_embd, config.n_embd)
-        self.query = nn.Linear(config.n_embd, config.n_embd)
-        self.value = nn.Linear(config.n_embd, config.n_embd)
-        # regularization
-        self.attn_drop = nn.Dropout(config.attn_pdrop)
-        self.resid_drop = nn.Dropout(config.resid_pdrop)
-        # output projection
-        self.proj = nn.Linear(config.n_embd, config.n_embd)
-        # causal mask to ensure that attention is only applied to the left in the input sequence
-        mask = torch.tril(torch.ones(config.block_size,
-                                     config.block_size))
-        if hasattr(config, "n_unmasked"):
-            mask[:config.n_unmasked, :config.n_unmasked] = 1
-        self.register_buffer("mask", mask.view(1, 1, config.block_size, config.block_size))
-        self.n_head = config.n_head
-
-    def forward(self, x, layer_past=None):
-        B, T, C = x.size()
-
-        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
-        k = self.key(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
-        q = self.query(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
-        v = self.value(x).view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
-
-        present = torch.stack((k, v))
-        if layer_past is not None:
-            past_key, past_value = layer_past
-            k = torch.cat((past_key, k), dim=-2)
-            v = torch.cat((past_value, v), dim=-2)
-
-        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
-        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
-        if layer_past is None:
-            att = att.masked_fill(self.mask[:,:,:T,:T] == 0, float('-inf'))
-
-        att = F.softmax(att, dim=-1)
-        att = self.attn_drop(att)
-        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
-        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
-
-        # output projection
-        y = self.resid_drop(self.proj(y))
-        return y, present   # TODO: check that this does not break anything
-
-
-class Block(nn.Module):
-    """ an unassuming Transformer block """
-    def __init__(self, config):
-        super().__init__()
-        self.ln1 = nn.LayerNorm(config.n_embd)
-        self.ln2 = nn.LayerNorm(config.n_embd)
-        self.attn = CausalSelfAttention(config)
-        self.mlp = nn.Sequential(
-            nn.Linear(config.n_embd, 4 * config.n_embd),
-            nn.GELU(),  # nice
-            nn.Linear(4 * config.n_embd, config.n_embd),
-            nn.Dropout(config.resid_pdrop),
-        )
-
-    def forward(self, x, layer_past=None, return_present=False):
-        # TODO: check that training still works
-        if return_present: assert not self.training
-        # layer past: tuple of length two with B, nh, T, hs
-        attn, present = self.attn(self.ln1(x), layer_past=layer_past)
-
-        x = x + attn
-        x = x + self.mlp(self.ln2(x))
-        if layer_past is not None or return_present:
-            return x, present
-        return x
-
-
-class GPT(nn.Module):
-    """  the full GPT language model, with a context size of block_size """
-    def __init__(self, vocab_size, block_size, n_layer=12, n_head=8, n_embd=256,
-                 embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0):
-        super().__init__()
-        config = GPTConfig(vocab_size=vocab_size, block_size=block_size,
-                           embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop,
-                           n_layer=n_layer, n_head=n_head, n_embd=n_embd,
-                           n_unmasked=n_unmasked)
-        # input embedding stem
-        self.tok_emb = nn.Embedding(config.vocab_size, config.n_embd)
-        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
-        self.drop = nn.Dropout(config.embd_pdrop)
-        # transformer
-        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
-        # decoder head
-        self.ln_f = nn.LayerNorm(config.n_embd)
-        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        self.block_size = config.block_size
-        self.apply(self._init_weights)
-        self.config = config
-        logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters()))
-
-    def get_block_size(self):
-        return self.block_size
-
-    def _init_weights(self, module):
-        if isinstance(module, (nn.Linear, nn.Embedding)):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if isinstance(module, nn.Linear) and module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def forward(self, idx, embeddings=None, targets=None):
-        # forward the GPT model
-        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
-
-        if embeddings is not None: # prepend explicit embeddings
-            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
-
-        t = token_embeddings.shape[1]
-        assert t <= self.block_size, "Cannot forward, model block size is exhausted."
-        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
-        x = self.drop(token_embeddings + position_embeddings)
-        x = self.blocks(x)
-        x = self.ln_f(x)
-        logits = self.head(x)
-
-        # if we are given some desired targets also calculate the loss
-        loss = None
-        if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
-
-        return logits, loss
-
-    def forward_with_past(self, idx, embeddings=None, targets=None, past=None, past_length=None):
-        # inference only
-        assert not self.training
-        token_embeddings = self.tok_emb(idx)    # each index maps to a (learnable) vector
-        if embeddings is not None:              # prepend explicit embeddings
-            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
-
-        if past is not None:
-            assert past_length is not None
-            past = torch.cat(past, dim=-2)   # n_layer, 2, b, nh, len_past, dim_head
-            past_shape = list(past.shape)
-            expected_shape = [self.config.n_layer, 2, idx.shape[0], self.config.n_head, past_length, self.config.n_embd//self.config.n_head]
-            assert past_shape == expected_shape, f"{past_shape} =/= {expected_shape}"
-            position_embeddings = self.pos_emb[:, past_length, :]  # each position maps to a (learnable) vector
-        else:
-            position_embeddings = self.pos_emb[:, :token_embeddings.shape[1], :]
-
-        x = self.drop(token_embeddings + position_embeddings)
-        presents = []  # accumulate over layers
-        for i, block in enumerate(self.blocks):
-            x, present = block(x, layer_past=past[i, ...] if past is not None else None, return_present=True)
-            presents.append(present)
-
-        x = self.ln_f(x)
-        logits = self.head(x)
-        # if we are given some desired targets also calculate the loss
-        loss = None
-        if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
-
-        return logits, loss, torch.stack(presents)  # _, _, n_layer, 2, b, nh, 1, dim_head
-
-
-class DummyGPT(nn.Module):
-    # for debugging
-    def __init__(self, add_value=1):
-        super().__init__()
-        self.add_value = add_value
-
-    def forward(self, idx):
-        return idx + self.add_value, None
-
-
-class CodeGPT(nn.Module):
-    """Takes in semi-embeddings"""
-    def __init__(self, vocab_size, block_size, in_channels, n_layer=12, n_head=8, n_embd=256,
-                 embd_pdrop=0., resid_pdrop=0., attn_pdrop=0., n_unmasked=0):
-        super().__init__()
-        config = GPTConfig(vocab_size=vocab_size, block_size=block_size,
-                           embd_pdrop=embd_pdrop, resid_pdrop=resid_pdrop, attn_pdrop=attn_pdrop,
-                           n_layer=n_layer, n_head=n_head, n_embd=n_embd,
-                           n_unmasked=n_unmasked)
-        # input embedding stem
-        self.tok_emb = nn.Linear(in_channels, config.n_embd)
-        self.pos_emb = nn.Parameter(torch.zeros(1, config.block_size, config.n_embd))
-        self.drop = nn.Dropout(config.embd_pdrop)
-        # transformer
-        self.blocks = nn.Sequential(*[Block(config) for _ in range(config.n_layer)])
-        # decoder head
-        self.ln_f = nn.LayerNorm(config.n_embd)
-        self.head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
-        self.block_size = config.block_size
-        self.apply(self._init_weights)
-        self.config = config
-        logger.info("number of parameters: %e", sum(p.numel() for p in self.parameters()))
-
-    def get_block_size(self):
-        return self.block_size
-
-    def _init_weights(self, module):
-        if isinstance(module, (nn.Linear, nn.Embedding)):
-            module.weight.data.normal_(mean=0.0, std=0.02)
-            if isinstance(module, nn.Linear) and module.bias is not None:
-                module.bias.data.zero_()
-        elif isinstance(module, nn.LayerNorm):
-            module.bias.data.zero_()
-            module.weight.data.fill_(1.0)
-
-    def forward(self, idx, embeddings=None, targets=None):
-        # forward the GPT model
-        token_embeddings = self.tok_emb(idx) # each index maps to a (learnable) vector
-
-        if embeddings is not None: # prepend explicit embeddings
-            token_embeddings = torch.cat((embeddings, token_embeddings), dim=1)
-
-        t = token_embeddings.shape[1]
-        assert t <= self.block_size, "Cannot forward, model block size is exhausted."
-        position_embeddings = self.pos_emb[:, :t, :] # each position maps to a (learnable) vector
-        x = self.drop(token_embeddings + position_embeddings)
-        x = self.blocks(x)
-        x = self.taming_cinln_f(x)
-        logits = self.head(x)
-
-        # if we are given some desired targets also calculate the loss
-        loss = None
-        if targets is not None:
-            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
-
-        return logits, loss
-
-
-
-#### sampling utils
-
-def top_k_logits(logits, k):
-    v, ix = torch.topk(logits, k)
-    out = logits.clone()
-    out[out < v[:, [-1]]] = -float('Inf')
-    return out
-
-@torch.no_grad()
-def sample(model, x, steps, temperature=1.0, sample=False, top_k=None):
-    """
-    take a conditioning sequence of indices in x (of shape (b,t)) and predict the next token in
-    the sequence, feeding the predictions back into the model each time. Clearly the sampling
-    has quadratic complexity unlike an RNN that is only linear, and has a finite context window
-    of block_size, unlike an RNN that has an infinite context window.
-    """
-    block_size = model.get_block_size()
-    model.eval()
-    for k in range(steps):
-        x_cond = x if x.size(1) <= block_size else x[:, -block_size:]  # crop context if needed
-        logits, _ = model(x_cond)
-        # pluck the logits at the final step and scale by temperature
-        logits = logits[:, -1, :] / temperature
-        # optionally crop probabilities to only the top k options
-        if top_k is not None:
-            logits = top_k_logits(logits, top_k)
-        # apply softmax to convert to probabilities
-        probs = F.softmax(logits, dim=-1)
-        # sample from the distribution or take the most likely
-        if sample:
-            ix = torch.multinomial(probs, num_samples=1)
-        else:
-            _, ix = torch.topk(probs, k=1, dim=-1)
-        # append to the sequence and continue
-        x = torch.cat((x, ix), dim=1)
-
-    return x
-
-
-@torch.no_grad()
-def sample_with_past(x, model, steps, temperature=1., sample_logits=True,
-                     top_k=None, top_p=None, callback=None):
-    # x is conditioning
-    sample = x
-    cond_len = x.shape[1]
-    past = None
-    for n in range(steps):
-        if callback is not None:
-            callback(n)
-        logits, _, present = model.forward_with_past(x, past=past, past_length=(n+cond_len-1))
-        if past is None:
-            past = [present]
-        else:
-            past.append(present)
-        logits = logits[:, -1, :] / temperature
-        if top_k is not None:
-            logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)
-
-        probs = F.softmax(logits, dim=-1)
-        if not sample_logits:
-            _, x = torch.topk(probs, k=1, dim=-1)
-        else:
-            x = torch.multinomial(probs, num_samples=1)
-        # append to the sequence and continue
-        sample = torch.cat((sample, x), dim=1)
-    del past
-    sample = sample[:, cond_len:]  # cut conditioning off
-    return sample
-
-
-#### clustering utils
-
-class KMeans(nn.Module):
-    def __init__(self, ncluster=512, nc=3, niter=10):
-        super().__init__()
-        self.ncluster = ncluster
-        self.nc = nc
-        self.niter = niter
-        self.shape = (3,32,32)
-        self.register_buffer("C", torch.zeros(self.ncluster,nc))
-        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
-
-    def is_initialized(self):
-        return self.initialized.item() == 1
-
-    @torch.no_grad()
-    def initialize(self, x):
-        N, D = x.shape
-        assert D == self.nc, D
-        c = x[torch.randperm(N)[:self.ncluster]] # init clusters at random
-        for i in range(self.niter):
-            # assign all pixels to the closest codebook element
-            a = ((x[:, None, :] - c[None, :, :])**2).sum(-1).argmin(1)
-            # move each codebook element to be the mean of the pixels that assigned to it
-            c = torch.stack([x[a==k].mean(0) for k in range(self.ncluster)])
-            # re-assign any poorly positioned codebook elements
-            nanix = torch.any(torch.isnan(c), dim=1)
-            ndead = nanix.sum().item()
-            print('done step %d/%d, re-initialized %d dead clusters' % (i+1, self.niter, ndead))
-            c[nanix] = x[torch.randperm(N)[:ndead]] # re-init dead clusters
-
-        self.C.copy_(c)
-        self.initialized.fill_(1)
-
-
-    def forward(self, x, reverse=False, shape=None):
-        if not reverse:
-            # flatten
-            bs,c,h,w = x.shape
-            assert c == self.nc
-            x = x.reshape(bs,c,h*w,1)
-            C = self.C.permute(1,0)
-            C = C.reshape(1,c,1,self.ncluster)
-            a = ((x-C)**2).sum(1).argmin(-1) # bs, h*w indices
-            return a
-        else:
-            # flatten
-            bs, HW = x.shape
-            """
-            c = self.C.reshape( 1, self.nc,  1, self.ncluster)
-            c = c[bs*[0],:,:,:]
-            c = c[:,:,HW*[0],:]
-            x =      x.reshape(bs,       1, HW,             1)
-            x = x[:,3*[0],:,:]
-            x = torch.gather(c, dim=3, index=x)
-            """
-            x = self.C[x]
-            x = x.permute(0,2,1)
-            shape = shape if shape is not None else self.shape
-            x = x.reshape(bs, *shape)
-
-            return x

taming-transformers/taming/modules/transformer/permuter.py

@@ -1,248 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-
-
-class AbstractPermuter(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-    def forward(self, x, reverse=False):
-        raise NotImplementedError
-
-
-class Identity(AbstractPermuter):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x, reverse=False):
-        return x
-
-
-class Subsample(AbstractPermuter):
-    def __init__(self, H, W):
-        super().__init__()
-        C = 1
-        indices = np.arange(H*W).reshape(C,H,W)
-        while min(H, W) > 1:
-            indices = indices.reshape(C,H//2,2,W//2,2)
-            indices = indices.transpose(0,2,4,1,3)
-            indices = indices.reshape(C*4,H//2, W//2)
-            H = H//2
-            W = W//2
-            C = C*4
-        assert H == W == 1
-        idx = torch.tensor(indices.ravel())
-        self.register_buffer('forward_shuffle_idx',
-                             nn.Parameter(idx, requires_grad=False))
-        self.register_buffer('backward_shuffle_idx',
-                             nn.Parameter(torch.argsort(idx), requires_grad=False))
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-def mortonify(i, j):
-    """(i,j) index to linear morton code"""
-    i = np.uint64(i)
-    j = np.uint64(j)
-
-    z = np.uint(0)
-
-    for pos in range(32):
-        z = (z |
-             ((j & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos)) |
-             ((i & (np.uint64(1) << np.uint64(pos))) << np.uint64(pos+1))
-             )
-    return z
-
-
-class ZCurve(AbstractPermuter):
-    def __init__(self, H, W):
-        super().__init__()
-        reverseidx = [np.int64(mortonify(i,j)) for i in range(H) for j in range(W)]
-        idx = np.argsort(reverseidx)
-        idx = torch.tensor(idx)
-        reverseidx = torch.tensor(reverseidx)
-        self.register_buffer('forward_shuffle_idx',
-                             idx)
-        self.register_buffer('backward_shuffle_idx',
-                             reverseidx)
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-class SpiralOut(AbstractPermuter):
-    def __init__(self, H, W):
-        super().__init__()
-        assert H == W
-        size = W
-        indices = np.arange(size*size).reshape(size,size)
-
-        i0 = size//2
-        j0 = size//2-1
-
-        i = i0
-        j = j0
-
-        idx = [indices[i0, j0]]
-        step_mult = 0
-        for c in range(1, size//2+1):
-            step_mult += 1
-            # steps left
-            for k in range(step_mult):
-                i = i - 1
-                j = j
-                idx.append(indices[i, j])
-
-            # step down
-            for k in range(step_mult):
-                i = i
-                j = j + 1
-                idx.append(indices[i, j])
-
-            step_mult += 1
-            if c < size//2:
-                # step right
-                for k in range(step_mult):
-                    i = i + 1
-                    j = j
-                    idx.append(indices[i, j])
-
-                # step up
-                for k in range(step_mult):
-                    i = i
-                    j = j - 1
-                    idx.append(indices[i, j])
-            else:
-                # end reached
-                for k in range(step_mult-1):
-                    i = i + 1
-                    idx.append(indices[i, j])
-
-        assert len(idx) == size*size
-        idx = torch.tensor(idx)
-        self.register_buffer('forward_shuffle_idx', idx)
-        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-class SpiralIn(AbstractPermuter):
-    def __init__(self, H, W):
-        super().__init__()
-        assert H == W
-        size = W
-        indices = np.arange(size*size).reshape(size,size)
-
-        i0 = size//2
-        j0 = size//2-1
-
-        i = i0
-        j = j0
-
-        idx = [indices[i0, j0]]
-        step_mult = 0
-        for c in range(1, size//2+1):
-            step_mult += 1
-            # steps left
-            for k in range(step_mult):
-                i = i - 1
-                j = j
-                idx.append(indices[i, j])
-
-            # step down
-            for k in range(step_mult):
-                i = i
-                j = j + 1
-                idx.append(indices[i, j])
-
-            step_mult += 1
-            if c < size//2:
-                # step right
-                for k in range(step_mult):
-                    i = i + 1
-                    j = j
-                    idx.append(indices[i, j])
-
-                # step up
-                for k in range(step_mult):
-                    i = i
-                    j = j - 1
-                    idx.append(indices[i, j])
-            else:
-                # end reached
-                for k in range(step_mult-1):
-                    i = i + 1
-                    idx.append(indices[i, j])
-
-        assert len(idx) == size*size
-        idx = idx[::-1]
-        idx = torch.tensor(idx)
-        self.register_buffer('forward_shuffle_idx', idx)
-        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-class Random(nn.Module):
-    def __init__(self, H, W):
-        super().__init__()
-        indices = np.random.RandomState(1).permutation(H*W)
-        idx = torch.tensor(indices.ravel())
-        self.register_buffer('forward_shuffle_idx', idx)
-        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-class AlternateParsing(AbstractPermuter):
-    def __init__(self, H, W):
-        super().__init__()
-        indices = np.arange(W*H).reshape(H,W)
-        for i in range(1, H, 2):
-            indices[i, :] = indices[i, ::-1]
-        idx = indices.flatten()
-        assert len(idx) == H*W
-        idx = torch.tensor(idx)
-        self.register_buffer('forward_shuffle_idx', idx)
-        self.register_buffer('backward_shuffle_idx', torch.argsort(idx))
-
-    def forward(self, x, reverse=False):
-        if not reverse:
-            return x[:, self.forward_shuffle_idx]
-        else:
-            return x[:, self.backward_shuffle_idx]
-
-
-if __name__ == "__main__":
-    p0 = AlternateParsing(16, 16)
-    print(p0.forward_shuffle_idx)
-    print(p0.backward_shuffle_idx)
-
-    x = torch.randint(0, 768, size=(11, 256))
-    y = p0(x)
-    xre = p0(y, reverse=True)
-    assert torch.equal(x, xre)
-
-    p1 = SpiralOut(2, 2)
-    print(p1.forward_shuffle_idx)
-    print(p1.backward_shuffle_idx)

taming-transformers/taming/modules/util.py

@@ -1,130 +0,0 @@
-import torch
-import torch.nn as nn
-
-
-def count_params(model):
-    total_params = sum(p.numel() for p in model.parameters())
-    return total_params
-
-
-class ActNorm(nn.Module):
-    def __init__(self, num_features, logdet=False, affine=True,
-                 allow_reverse_init=False):
-        assert affine
-        super().__init__()
-        self.logdet = logdet
-        self.loc = nn.Parameter(torch.zeros(1, num_features, 1, 1))
-        self.scale = nn.Parameter(torch.ones(1, num_features, 1, 1))
-        self.allow_reverse_init = allow_reverse_init
-
-        self.register_buffer('initialized', torch.tensor(0, dtype=torch.uint8))
-
-    def initialize(self, input):
-        with torch.no_grad():
-            flatten = input.permute(1, 0, 2, 3).contiguous().view(input.shape[1], -1)
-            mean = (
-                flatten.mean(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-            std = (
-                flatten.std(1)
-                .unsqueeze(1)
-                .unsqueeze(2)
-                .unsqueeze(3)
-                .permute(1, 0, 2, 3)
-            )
-
-            self.loc.data.copy_(-mean)
-            self.scale.data.copy_(1 / (std + 1e-6))
-
-    def forward(self, input, reverse=False):
-        if reverse:
-            return self.reverse(input)
-        if len(input.shape) == 2:
-            input = input[:,:,None,None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        _, _, height, width = input.shape
-
-        if self.training and self.initialized.item() == 0:
-            self.initialize(input)
-            self.initialized.fill_(1)
-
-        h = self.scale * (input + self.loc)
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-
-        if self.logdet:
-            log_abs = torch.log(torch.abs(self.scale))
-            logdet = height*width*torch.sum(log_abs)
-            logdet = logdet * torch.ones(input.shape[0]).to(input)
-            return h, logdet
-
-        return h
-
-    def reverse(self, output):
-        if self.training and self.initialized.item() == 0:
-            if not self.allow_reverse_init:
-                raise RuntimeError(
-                    "Initializing ActNorm in reverse direction is "
-                    "disabled by default. Use allow_reverse_init=True to enable."
-                )
-            else:
-                self.initialize(output)
-                self.initialized.fill_(1)
-
-        if len(output.shape) == 2:
-            output = output[:,:,None,None]
-            squeeze = True
-        else:
-            squeeze = False
-
-        h = output / self.scale - self.loc
-
-        if squeeze:
-            h = h.squeeze(-1).squeeze(-1)
-        return h
-
-
-class AbstractEncoder(nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def encode(self, *args, **kwargs):
-        raise NotImplementedError
-
-
-class Labelator(AbstractEncoder):
-    """Net2Net Interface for Class-Conditional Model"""
-    def __init__(self, n_classes, quantize_interface=True):
-        super().__init__()
-        self.n_classes = n_classes
-        self.quantize_interface = quantize_interface
-
-    def encode(self, c):
-        c = c[:,None]
-        if self.quantize_interface:
-            return c, None, [None, None, c.long()]
-        return c
-
-
-class SOSProvider(AbstractEncoder):
-    # for unconditional training
-    def __init__(self, sos_token, quantize_interface=True):
-        super().__init__()
-        self.sos_token = sos_token
-        self.quantize_interface = quantize_interface
-
-    def encode(self, x):
-        # get batch size from data and replicate sos_token
-        c = torch.ones(x.shape[0], 1)*self.sos_token
-        c = c.long().to(x.device)
-        if self.quantize_interface:
-            return c, None, [None, None, c]
-        return c

taming-transformers/taming/modules/vqvae/quantize.py

@@ -1,445 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from torch import einsum
-from einops import rearrange
-
-
-class VectorQuantizer(nn.Module):
-    """
-    see https://github.com/MishaLaskin/vqvae/blob/d761a999e2267766400dc646d82d3ac3657771d4/models/quantizer.py
-    ____________________________________________
-    Discretization bottleneck part of the VQ-VAE.
-    Inputs:
-    - n_e : number of embeddings
-    - e_dim : dimension of embedding
-    - beta : commitment cost used in loss term, beta * ||z_e(x)-sg[e]||^2
-    _____________________________________________
-    """
-
-    # NOTE: this class contains a bug regarding beta; see VectorQuantizer2 for
-    # a fix and use legacy=False to apply that fix. VectorQuantizer2 can be
-    # used wherever VectorQuantizer has been used before and is additionally
-    # more efficient.
-    def __init__(self, n_e, e_dim, beta):
-        super(VectorQuantizer, self).__init__()
-        self.n_e = n_e
-        self.e_dim = e_dim
-        self.beta = beta
-
-        self.embedding = nn.Embedding(self.n_e, self.e_dim)
-        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
-
-    def forward(self, z):
-        """
-        Inputs the output of the encoder network z and maps it to a discrete
-        one-hot vector that is the index of the closest embedding vector e_j
-        z (continuous) -> z_q (discrete)
-        z.shape = (batch, channel, height, width)
-        quantization pipeline:
-            1. get encoder input (B,C,H,W)
-            2. flatten input to (B*H*W,C)
-        """
-        # reshape z -> (batch, height, width, channel) and flatten
-        z = z.permute(0, 2, 3, 1).contiguous()
-        z_flattened = z.view(-1, self.e_dim)
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-
-        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
-            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
-            torch.matmul(z_flattened, self.embedding.weight.t())
-
-        ## could possible replace this here
-        # #\start...
-        # find closest encodings
-        min_encoding_indices = torch.argmin(d, dim=1).unsqueeze(1)
-
-        min_encodings = torch.zeros(
-            min_encoding_indices.shape[0], self.n_e).to(z)
-        min_encodings.scatter_(1, min_encoding_indices, 1)
-
-        # dtype min encodings: torch.float32
-        # min_encodings shape: torch.Size([2048, 512])
-        # min_encoding_indices.shape: torch.Size([2048, 1])
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings, self.embedding.weight).view(z.shape)
-        #.........\end
-
-        # with:
-        # .........\start
-        #min_encoding_indices = torch.argmin(d, dim=1)
-        #z_q = self.embedding(min_encoding_indices)
-        # ......\end......... (TODO)
-
-        # compute loss for embedding
-        loss = torch.mean((z_q.detach()-z)**2) + self.beta * \
-            torch.mean((z_q - z.detach()) ** 2)
-
-        # preserve gradients
-        z_q = z + (z_q - z).detach()
-
-        # perplexity
-        e_mean = torch.mean(min_encodings, dim=0)
-        perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + 1e-10)))
-
-        # reshape back to match original input shape
-        z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
-
-    def get_codebook_entry(self, indices, shape):
-        # shape specifying (batch, height, width, channel)
-        # TODO: check for more easy handling with nn.Embedding
-        min_encodings = torch.zeros(indices.shape[0], self.n_e).to(indices)
-        min_encodings.scatter_(1, indices[:,None], 1)
-
-        # get quantized latent vectors
-        z_q = torch.matmul(min_encodings.float(), self.embedding.weight)
-
-        if shape is not None:
-            z_q = z_q.view(shape)
-
-            # reshape back to match original input shape
-            z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q
-
-
-class GumbelQuantize(nn.Module):
-    """
-    credit to @karpathy: https://github.com/karpathy/deep-vector-quantization/blob/main/model.py (thanks!)
-    Gumbel Softmax trick quantizer
-    Categorical Reparameterization with Gumbel-Softmax, Jang et al. 2016
-    https://arxiv.org/abs/1611.01144
-    """
-    def __init__(self, num_hiddens, embedding_dim, n_embed, straight_through=True,
-                 kl_weight=5e-4, temp_init=1.0, use_vqinterface=True,
-                 remap=None, unknown_index="random"):
-        super().__init__()
-
-        self.embedding_dim = embedding_dim
-        self.n_embed = n_embed
-
-        self.straight_through = straight_through
-        self.temperature = temp_init
-        self.kl_weight = kl_weight
-
-        self.proj = nn.Conv2d(num_hiddens, n_embed, 1)
-        self.embed = nn.Embedding(n_embed, embedding_dim)
-
-        self.use_vqinterface = use_vqinterface
-
-        self.remap = remap
-        if self.remap is not None:
-            self.register_buffer("used", torch.tensor(np.load(self.remap)))
-            self.re_embed = self.used.shape[0]
-            self.unknown_index = unknown_index # "random" or "extra" or integer
-            if self.unknown_index == "extra":
-                self.unknown_index = self.re_embed
-                self.re_embed = self.re_embed+1
-            print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
-                  f"Using {self.unknown_index} for unknown indices.")
-        else:
-            self.re_embed = n_embed
-
-    def remap_to_used(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        match = (inds[:,:,None]==used[None,None,...]).long()
-        new = match.argmax(-1)
-        unknown = match.sum(2)<1
-        if self.unknown_index == "random":
-            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
-        else:
-            new[unknown] = self.unknown_index
-        return new.reshape(ishape)
-
-    def unmap_to_all(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        if self.re_embed > self.used.shape[0]: # extra token
-            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
-        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
-        return back.reshape(ishape)
-
-    def forward(self, z, temp=None, return_logits=False):
-        # force hard = True when we are in eval mode, as we must quantize. actually, always true seems to work
-        hard = self.straight_through if self.training else True
-        temp = self.temperature if temp is None else temp
-
-        logits = self.proj(z)
-        if self.remap is not None:
-            # continue only with used logits
-            full_zeros = torch.zeros_like(logits)
-            logits = logits[:,self.used,...]
-
-        soft_one_hot = F.gumbel_softmax(logits, tau=temp, dim=1, hard=hard)
-        if self.remap is not None:
-            # go back to all entries but unused set to zero
-            full_zeros[:,self.used,...] = soft_one_hot
-            soft_one_hot = full_zeros
-        z_q = einsum('b n h w, n d -> b d h w', soft_one_hot, self.embed.weight)
-
-        # + kl divergence to the prior loss
-        qy = F.softmax(logits, dim=1)
-        diff = self.kl_weight * torch.sum(qy * torch.log(qy * self.n_embed + 1e-10), dim=1).mean()
-
-        ind = soft_one_hot.argmax(dim=1)
-        if self.remap is not None:
-            ind = self.remap_to_used(ind)
-        if self.use_vqinterface:
-            if return_logits:
-                return z_q, diff, (None, None, ind), logits
-            return z_q, diff, (None, None, ind)
-        return z_q, diff, ind
-
-    def get_codebook_entry(self, indices, shape):
-        b, h, w, c = shape
-        assert b*h*w == indices.shape[0]
-        indices = rearrange(indices, '(b h w) -> b h w', b=b, h=h, w=w)
-        if self.remap is not None:
-            indices = self.unmap_to_all(indices)
-        one_hot = F.one_hot(indices, num_classes=self.n_embed).permute(0, 3, 1, 2).float()
-        z_q = einsum('b n h w, n d -> b d h w', one_hot, self.embed.weight)
-        return z_q
-
-
-class VectorQuantizer2(nn.Module):
-    """
-    Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly
-    avoids costly matrix multiplications and allows for post-hoc remapping of indices.
-    """
-    # NOTE: due to a bug the beta term was applied to the wrong term. for
-    # backwards compatibility we use the buggy version by default, but you can
-    # specify legacy=False to fix it.
-    def __init__(self, n_e, e_dim, beta, remap=None, unknown_index="random",
-                 sane_index_shape=False, legacy=True):
-        super().__init__()
-        self.n_e = n_e
-        self.e_dim = e_dim
-        self.beta = beta
-        self.legacy = legacy
-
-        self.embedding = nn.Embedding(self.n_e, self.e_dim)
-        self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
-
-        self.remap = remap
-        if self.remap is not None:
-            self.register_buffer("used", torch.tensor(np.load(self.remap)))
-            self.re_embed = self.used.shape[0]
-            self.unknown_index = unknown_index # "random" or "extra" or integer
-            if self.unknown_index == "extra":
-                self.unknown_index = self.re_embed
-                self.re_embed = self.re_embed+1
-            print(f"Remapping {self.n_e} indices to {self.re_embed} indices. "
-                  f"Using {self.unknown_index} for unknown indices.")
-        else:
-            self.re_embed = n_e
-
-        self.sane_index_shape = sane_index_shape
-
-    def remap_to_used(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        match = (inds[:,:,None]==used[None,None,...]).long()
-        new = match.argmax(-1)
-        unknown = match.sum(2)<1
-        if self.unknown_index == "random":
-            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
-        else:
-            new[unknown] = self.unknown_index
-        return new.reshape(ishape)
-
-    def unmap_to_all(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        if self.re_embed > self.used.shape[0]: # extra token
-            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
-        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
-        return back.reshape(ishape)
-
-    def forward(self, z, temp=None, rescale_logits=False, return_logits=False):
-        assert temp is None or temp==1.0, "Only for interface compatible with Gumbel"
-        assert rescale_logits==False, "Only for interface compatible with Gumbel"
-        assert return_logits==False, "Only for interface compatible with Gumbel"
-        # reshape z -> (batch, height, width, channel) and flatten
-        z = rearrange(z, 'b c h w -> b h w c').contiguous()
-        z_flattened = z.view(-1, self.e_dim)
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-
-        d = torch.sum(z_flattened ** 2, dim=1, keepdim=True) + \
-            torch.sum(self.embedding.weight**2, dim=1) - 2 * \
-            torch.einsum('bd,dn->bn', z_flattened, rearrange(self.embedding.weight, 'n d -> d n'))
-
-        min_encoding_indices = torch.argmin(d, dim=1)
-        z_q = self.embedding(min_encoding_indices).view(z.shape)
-        perplexity = None
-        min_encodings = None
-
-        # compute loss for embedding
-        if not self.legacy:
-            loss = self.beta * torch.mean((z_q.detach()-z)**2) + \
-                   torch.mean((z_q - z.detach()) ** 2)
-        else:
-            loss = torch.mean((z_q.detach()-z)**2) + self.beta * \
-                   torch.mean((z_q - z.detach()) ** 2)
-
-        # preserve gradients
-        z_q = z + (z_q - z).detach()
-
-        # reshape back to match original input shape
-        z_q = rearrange(z_q, 'b h w c -> b c h w').contiguous()
-
-        if self.remap is not None:
-            min_encoding_indices = min_encoding_indices.reshape(z.shape[0],-1) # add batch axis
-            min_encoding_indices = self.remap_to_used(min_encoding_indices)
-            min_encoding_indices = min_encoding_indices.reshape(-1,1) # flatten
-
-        if self.sane_index_shape:
-            min_encoding_indices = min_encoding_indices.reshape(
-                z_q.shape[0], z_q.shape[2], z_q.shape[3])
-
-        return z_q, loss, (perplexity, min_encodings, min_encoding_indices)
-
-    def get_codebook_entry(self, indices, shape):
-        # shape specifying (batch, height, width, channel)
-        if self.remap is not None:
-            indices = indices.reshape(shape[0],-1) # add batch axis
-            indices = self.unmap_to_all(indices)
-            indices = indices.reshape(-1) # flatten again
-
-        # get quantized latent vectors
-        z_q = self.embedding(indices)
-
-        if shape is not None:
-            z_q = z_q.view(shape)
-            # reshape back to match original input shape
-            z_q = z_q.permute(0, 3, 1, 2).contiguous()
-
-        return z_q
-
-class EmbeddingEMA(nn.Module):
-    def __init__(self, num_tokens, codebook_dim, decay=0.99, eps=1e-5):
-        super().__init__()
-        self.decay = decay
-        self.eps = eps        
-        weight = torch.randn(num_tokens, codebook_dim)
-        self.weight = nn.Parameter(weight, requires_grad = False)
-        self.cluster_size = nn.Parameter(torch.zeros(num_tokens), requires_grad = False)
-        self.embed_avg = nn.Parameter(weight.clone(), requires_grad = False)
-        self.update = True
-
-    def forward(self, embed_id):
-        return F.embedding(embed_id, self.weight)
-
-    def cluster_size_ema_update(self, new_cluster_size):
-        self.cluster_size.data.mul_(self.decay).add_(new_cluster_size, alpha=1 - self.decay)
-
-    def embed_avg_ema_update(self, new_embed_avg): 
-        self.embed_avg.data.mul_(self.decay).add_(new_embed_avg, alpha=1 - self.decay)
-
-    def weight_update(self, num_tokens):
-        n = self.cluster_size.sum()
-        smoothed_cluster_size = (
-                (self.cluster_size + self.eps) / (n + num_tokens * self.eps) * n
-            )
-        #normalize embedding average with smoothed cluster size
-        embed_normalized = self.embed_avg / smoothed_cluster_size.unsqueeze(1)
-        self.weight.data.copy_(embed_normalized)   
-
-
-class EMAVectorQuantizer(nn.Module):
-    def __init__(self, n_embed, embedding_dim, beta, decay=0.99, eps=1e-5,
-                remap=None, unknown_index="random"):
-        super().__init__()
-        self.codebook_dim = codebook_dim
-        self.num_tokens = num_tokens
-        self.beta = beta
-        self.embedding = EmbeddingEMA(self.num_tokens, self.codebook_dim, decay, eps)
-
-        self.remap = remap
-        if self.remap is not None:
-            self.register_buffer("used", torch.tensor(np.load(self.remap)))
-            self.re_embed = self.used.shape[0]
-            self.unknown_index = unknown_index # "random" or "extra" or integer
-            if self.unknown_index == "extra":
-                self.unknown_index = self.re_embed
-                self.re_embed = self.re_embed+1
-            print(f"Remapping {self.n_embed} indices to {self.re_embed} indices. "
-                  f"Using {self.unknown_index} for unknown indices.")
-        else:
-            self.re_embed = n_embed
-
-    def remap_to_used(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        match = (inds[:,:,None]==used[None,None,...]).long()
-        new = match.argmax(-1)
-        unknown = match.sum(2)<1
-        if self.unknown_index == "random":
-            new[unknown]=torch.randint(0,self.re_embed,size=new[unknown].shape).to(device=new.device)
-        else:
-            new[unknown] = self.unknown_index
-        return new.reshape(ishape)
-
-    def unmap_to_all(self, inds):
-        ishape = inds.shape
-        assert len(ishape)>1
-        inds = inds.reshape(ishape[0],-1)
-        used = self.used.to(inds)
-        if self.re_embed > self.used.shape[0]: # extra token
-            inds[inds>=self.used.shape[0]] = 0 # simply set to zero
-        back=torch.gather(used[None,:][inds.shape[0]*[0],:], 1, inds)
-        return back.reshape(ishape)
-
-    def forward(self, z):
-        # reshape z -> (batch, height, width, channel) and flatten
-        #z, 'b c h w -> b h w c'
-        z = rearrange(z, 'b c h w -> b h w c')
-        z_flattened = z.reshape(-1, self.codebook_dim)
-        
-        # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
-        d = z_flattened.pow(2).sum(dim=1, keepdim=True) + \
-            self.embedding.weight.pow(2).sum(dim=1) - 2 * \
-            torch.einsum('bd,nd->bn', z_flattened, self.embedding.weight) # 'n d -> d n'
-
-
-        encoding_indices = torch.argmin(d, dim=1)
-
-        z_q = self.embedding(encoding_indices).view(z.shape)
-        encodings = F.one_hot(encoding_indices, self.num_tokens).type(z.dtype)     
-        avg_probs = torch.mean(encodings, dim=0)
-        perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10)))
-
-        if self.training and self.embedding.update:
-            #EMA cluster size
-            encodings_sum = encodings.sum(0)            
-            self.embedding.cluster_size_ema_update(encodings_sum)
-            #EMA embedding average
-            embed_sum = encodings.transpose(0,1) @ z_flattened            
-            self.embedding.embed_avg_ema_update(embed_sum)
-            #normalize embed_avg and update weight
-            self.embedding.weight_update(self.num_tokens)
-
-        # compute loss for embedding
-        loss = self.beta * F.mse_loss(z_q.detach(), z) 
-
-        # preserve gradients
-        z_q = z + (z_q - z).detach()
-
-        # reshape back to match original input shape
-        #z_q, 'b h w c -> b c h w'
-        z_q = rearrange(z_q, 'b h w c -> b c h w')
-        return z_q, loss, (perplexity, encodings, encoding_indices)

taming-transformers/taming/util.py

@@ -1,157 +0,0 @@
-import os, hashlib
-import requests
-from tqdm import tqdm
-
-URL_MAP = {
-    "vgg_lpips": "https://heibox.uni-heidelberg.de/f/607503859c864bc1b30b/?dl=1"
-}
-
-CKPT_MAP = {
-    "vgg_lpips": "vgg.pth"
-}
-
-MD5_MAP = {
-    "vgg_lpips": "d507d7349b931f0638a25a48a722f98a"
-}
-
-
-def download(url, local_path, chunk_size=1024):
-    os.makedirs(os.path.split(local_path)[0], exist_ok=True)
-    with requests.get(url, stream=True) as r:
-        total_size = int(r.headers.get("content-length", 0))
-        with tqdm(total=total_size, unit="B", unit_scale=True) as pbar:
-            with open(local_path, "wb") as f:
-                for data in r.iter_content(chunk_size=chunk_size):
-                    if data:
-                        f.write(data)
-                        pbar.update(chunk_size)
-
-
-def md5_hash(path):
-    with open(path, "rb") as f:
-        content = f.read()
-    return hashlib.md5(content).hexdigest()
-
-
-def get_ckpt_path(name, root, check=False):
-    assert name in URL_MAP
-    path = os.path.join(root, CKPT_MAP[name])
-    if not os.path.exists(path) or (check and not md5_hash(path) == MD5_MAP[name]):
-        print("Downloading {} model from {} to {}".format(name, URL_MAP[name], path))
-        download(URL_MAP[name], path)
-        md5 = md5_hash(path)
-        assert md5 == MD5_MAP[name], md5
-    return path
-
-
-class KeyNotFoundError(Exception):
-    def __init__(self, cause, keys=None, visited=None):
-        self.cause = cause
-        self.keys = keys
-        self.visited = visited
-        messages = list()
-        if keys is not None:
-            messages.append("Key not found: {}".format(keys))
-        if visited is not None:
-            messages.append("Visited: {}".format(visited))
-        messages.append("Cause:\n{}".format(cause))
-        message = "\n".join(messages)
-        super().__init__(message)
-
-
-def retrieve(
-    list_or_dict, key, splitval="/", default=None, expand=True, pass_success=False
-):
-    """Given a nested list or dict return the desired value at key expanding
-    callable nodes if necessary and :attr:`expand` is ``True``. The expansion
-    is done in-place.
-
-    Parameters
-    ----------
-        list_or_dict : list or dict
-            Possibly nested list or dictionary.
-        key : str
-            key/to/value, path like string describing all keys necessary to
-            consider to get to the desired value. List indices can also be
-            passed here.
-        splitval : str
-            String that defines the delimiter between keys of the
-            different depth levels in `key`.
-        default : obj
-            Value returned if :attr:`key` is not found.
-        expand : bool
-            Whether to expand callable nodes on the path or not.
-
-    Returns
-    -------
-        The desired value or if :attr:`default` is not ``None`` and the
-        :attr:`key` is not found returns ``default``.
-
-    Raises
-    ------
-        Exception if ``key`` not in ``list_or_dict`` and :attr:`default` is
-        ``None``.
-    """
-
-    keys = key.split(splitval)
-
-    success = True
-    try:
-        visited = []
-        parent = None
-        last_key = None
-        for key in keys:
-            if callable(list_or_dict):
-                if not expand:
-                    raise KeyNotFoundError(
-                        ValueError(
-                            "Trying to get past callable node with expand=False."
-                        ),
-                        keys=keys,
-                        visited=visited,
-                    )
-                list_or_dict = list_or_dict()
-                parent[last_key] = list_or_dict
-
-            last_key = key
-            parent = list_or_dict
-
-            try:
-                if isinstance(list_or_dict, dict):
-                    list_or_dict = list_or_dict[key]
-                else:
-                    list_or_dict = list_or_dict[int(key)]
-            except (KeyError, IndexError, ValueError) as e:
-                raise KeyNotFoundError(e, keys=keys, visited=visited)
-
-            visited += [key]
-        # final expansion of retrieved value
-        if expand and callable(list_or_dict):
-            list_or_dict = list_or_dict()
-            parent[last_key] = list_or_dict
-    except KeyNotFoundError as e:
-        if default is None:
-            raise e
-        else:
-            list_or_dict = default
-            success = False
-
-    if not pass_success:
-        return list_or_dict
-    else:
-        return list_or_dict, success
-
-
-if __name__ == "__main__":
-    config = {"keya": "a",
-              "keyb": "b",
-              "keyc":
-                  {"cc1": 1,
-                   "cc2": 2,
-                   }
-              }
-    from omegaconf import OmegaConf
-    config = OmegaConf.create(config)
-    print(config)
-    retrieve(config, "keya")
-

taming-transformers/taming_transformers.egg-info/PKG-INFO

@@ -1,10 +0,0 @@
-Metadata-Version: 2.1
-Name: taming-transformers
-Version: 0.0.1
-Summary: Taming Transformers for High-Resolution Image Synthesis
-Home-page: UNKNOWN
-License: UNKNOWN
-Platform: UNKNOWN
-
-UNKNOWN
-

taming-transformers/taming_transformers.egg-info/SOURCES.txt

@@ -1,7 +0,0 @@
-README.md
-setup.py
-taming_transformers.egg-info/PKG-INFO
-taming_transformers.egg-info/SOURCES.txt
-taming_transformers.egg-info/dependency_links.txt
-taming_transformers.egg-info/requires.txt
-taming_transformers.egg-info/top_level.txt

taming-transformers/taming_transformers.egg-info/dependency_links.txt

@@ -1 +0,0 @@
-

taming-transformers/taming_transformers.egg-info/requires.txt

@@ -1,3 +0,0 @@
-torch
-numpy
-tqdm

taming-transformers/taming_transformers.egg-info/top_level.txt

@@ -1 +0,0 @@
-