dalle-pytorch

To support Chinese text, use the --chinese flag during both training and generation. This utilizes a pretrained Chinese tokenizer from Huggingface.

# Training
python train_dalle.py --chinese --image_text_folder ./path/to/data

# Generation
python generate.py --chinese --text '追老鼠的猫'

You can choose between different VAE architectures during training:

1. OpenAI's Pretrained VAE: Use by default by omitting the --vae_path flag.
2. Taming Transformer's VQGAN VAE: Highly recommended. Enable by adding the --taming flag.

To train the DALL-E model, use the train_dalle.py script.

Dataset Format: Provide a folder containing images and corresponding .txt files with the same name. The text files should contain multiple descriptions delimited by newlines. The script will randomly select one description per batch.

Commands:

• New training: python train_dalle.py --vae_path ./vae.pt --image_text_folder /path/to/data
• Resume training: python train_dalle.py --dalle_path ./dalle.pt --image_text_folder /path/to/data

To implement guidance without a classifier, you can randomly drop out the text condition during training and scale it during inference.

• During training: Set null_cond_prob (e.g., 0.2) in the dalle() call to define the probability of dropping the condition.
• During generation: Set cond_scale (e.g., 3.0) in the generate_images() call to increase conditioning strength.

The VAE is trained to compress images into discrete visual tokens. Use the train_vae.py script and provide an image folder. It is recommended to use wandb for experiment tracking (monitor reconstruction, temperature annealing, and codebook collapse). Model weights are saved to ./vae.pt.

$ pip install wandb
$ wandb login
$ python train_vae.py --image_folder /path/to/your/images

To use Microsoft Deepspeed's Sparse Attention with attn_types = ('full', 'sparse'), you must perform a specific installation:

1. Run the provided installation script: sh install_deepspeed.sh
2. Install triton version < 1.0 (specifically 0.4.2 is recommended): pip install triton==0.4.2

$ sh install_deepspeed.sh
$ pip install triton==0.4.2

You can replace the default tokenizer by providing a path to a BPE model file via the --bpe_path flag in both train_dalle.py and generate.py.

Requirement: You must use 0 as the padding during tokenization.

Workflow:

1. Install: pip install youtokentome
2. Create model: yttm bpe --vocab_size 8000 --data ./path/to/big/text/file.txt --model ./path/to/bpe.model
3. Use in training/generation: --bpe_path ./path/to/bpe.model

Replace python with deepspeed and add the --deepspeed flag to enable distributed training.

Precision Notes:

• For DeepSpeed < 0.3.16, you must pass the --fp16 flag to enable ZeRO optimizations.
• For Apex Automatic Mixed Precision (AMP), install DeepSpeed via the provided Dockerfile/script, install apex from source using sh install_apex.sh, and use the --amp flag with --distr_backend 'deepspeed'.

Use Docker to ensure correct PyTorch and CUDA versions.

To build the image: docker build -t dalle docker

To run in an interactive shell with GPU support and the current directory mounted to /workspace/dalle: docker run --gpus all -it --mount src="$(pwd)",target=/workspace/dalle,type=bind dalle:latest bash

Install the library using pip to get started with DALL-E implementation in PyTorch.

$ pip install dalle-pytorch