{"id":13564030,"url":"https://github.com/AranKomat/Diff-DALLE","last_synced_at":"2025-04-03T20:32:39.797Z","repository":{"id":113169478,"uuid":"380407448","full_name":"AranKomat/Diff-DALLE","owner":"AranKomat","description":null,"archived":false,"fork":false,"pushed_at":"2021-11-04T10:34:32.000Z","size":1407,"stargazers_count":64,"open_issues_count":1,"forks_count":8,"subscribers_count":11,"default_branch":"main","last_synced_at":"2025-03-24T19:45:15.851Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/AranKomat.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null}},"created_at":"2021-06-26T03:48:06.000Z","updated_at":"2024-02-08T20:01:30.000Z","dependencies_parsed_at":"2023-07-16T18:30:36.538Z","dependency_job_id":null,"html_url":"https://github.com/AranKomat/Diff-DALLE","commit_stats":null,"previous_names":[],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AranKomat%2FDiff-DALLE","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AranKomat%2FDiff-DALLE/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AranKomat%2FDiff-DALLE/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/AranKomat%2FDiff-DALLE/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/AranKomat","download_url":"https://codeload.github.com/AranKomat/Diff-DALLE/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":247075625,"owners_count":20879479,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2024-08-01T13:01:25.703Z","updated_at":"2025-04-03T20:32:39.474Z","avatar_url":"https://github.com/AranKomat.png","language":"Python","funding_links":[],"categories":["Python"],"sub_categories":[],"readme":"# Diff-DALLE (WIP)\n\nDiff-DALLE is DDPM + CLIP with 1.5B parameters for text-to-image generation. This repo allows both training and inference of this model. The links to pretrained weights and Colab notebook are attached.   \n\n-  Diff-DALLE consists of three generators and one classifier:\n\t- Generators are responsible for: \n\t\t- 64x64 text-to-image generation\n\t\t- 64x64 -\u003e 256x256 upscaling\n\t\t- 256x256 -\u003e 1024x1024 upscaling\n\t- Each generator consists of:\n\t\t- Transformer encoder for taking in the input text\n\t\t- DDPM U-Net for generating an image conditioned on the input text\n\t\t\t- Input embeddings are fed from the encoder to the U-Net via encoder-decoder attention.\n\t- CLIP classifier trained on noised images from scratch for guiding the sampling.\n\t\t- The image part of our CLIP uses the same architecture as in Guided Diffusion for better low-level guidance.\n- Pretrained model is trained on ~100M image-text pairs and ~40M high-resolution images with 32 A100 GPUs for ~6 weeks.\n\nWe have designed Diff-DALLE by starting from [Guided Diffusion](https://github.com/openai/guided-diffusion/) by Dhariwal et. al. Hence, its design, hyperparameter choice and code are heavily inspired by this. [Concurrent work](https://twitter.com/RiversHaveWings/status/1417629128124076032) by [Katherine Crowson](https://twitter.com/RiversHaveWings) on Guided Diffusion + CLIP is also worth noting (comparison below). \n\n#### Major steps to follow\n1.  Installation\n2. Data preparation\n3. Training generator (encoder + U-Net)\n4. Training classifier\n5. Sampling\n\n## Acknowledgment \nThanks to everyone who have helped out one way or another (listed alphabetically):\n* [Ben Wang](https://github.com/kingoflolz) for offering some valuable advices.\n* [CoreWeave](https://www.coreweave.com/) and [EleutherAI](https://www.eleuther.ai/) for providing their computational resources for training.\n* [Christoph Schuhmann](http://christoph-schuhmann.de/), [Richard Vencu](https://github.com/rvencu), [Romain Beaumont](https://rom1504.fr/) and many others in dalle-pytorch Discord server for building the Crawling@Home dataset, which was used along with our high-resolution dataset. Notably, Romain helped us to download the dataset.\n* [Shivanshu Purohit](https://github.com/ShivanshuPurohit) for securing and setting up the computational resources used for this project \n* [Spell](https://spell.ml/) for offering us a generous grant of $40,000 worth-compute from [Spell Open Research Grant](https://spell.ml/blog/spell-open-research-grant-YIMtSxEAACQAKm81) \n\n## Installation\n\nClone this repository and navigate to it in your terminal. Then run:\n\n```bash\npip install -e .\n```\n\nThis should install the `diff_dalle` python package that the scripts depend on.\n\n##### Caveat\n* Make sure to use the newest stable version of PyTorch (1.9.0), as older version slows down the training when used with Ampere (e.g. A100) architecture.\n\n## Preparing data\n\nCreate a directory with containing shards of webdataset consisting of images (and texts if applicable) and set `index_dir` to the path to the directory. \n* C@H dataset (400M CLIP-filtered image-text pairs) will be available at the-eye soon\n\t* Until then you can download the images at https://github.com/rom1504/img2dataset\n* ~35M high-resolution image dataset at https://the-eye.eu/eleuther_staging/yfcc2/\n\n## Training\n\n### General remarks\n* To generate with image resolution larger than 128, cascading training should be used. \n\t* For example, for 256 x 256  images, you need to train 64 x 64 generator + classifier as well as 256 x 256 upsampling generator (optionally + classifier).   \n\t* While a model with a higher-resolution hierarchy costs more FLOPS per image, you can mitigate the increased cost without substantial performance degradation by doing the following:\n\t\t* Use a smaller model  (e.g. `num_channels` = 256 for 64 x 64, while = 128 for 256 x 256)\n\t\t* Reduce `batch_size` (e.g. `batch_size` = 2048 for 64 x 64, while = 256 for 256 x 256)\n\t* This way, you can aim for spending about the same amount of computes on each hierarchy. \n\t* It is often more efficient to make the number of sampling steps larger for low-res hierarchy than high-res hierarchy.\n\n### Logs \u0026 checkpoints\n\nThe logs and checkpoints will be written to a logging directory determined by the `OPENAI_LOGDIR` environment variable (e.g. `export OPENAI_LOGDIR=/path/to/logdir`). If it is not set, then a temporary directory will be created in `/tmp`.\n\nThe training scripts below save checkpoints to `.pt` files in the logging directory. These checkpoints will have names like `ema_0.9999_200000.pt` and `model200000.pt`. You will likely want to sample from the EMA models, since those produce much better samples.\n\n### Generator\nTo train your model, you should first decide some hyperparameters. We will split up our hyperparameters into three groups: model architecture, diffusion process (for generator), and training flags. Here is an example:\n\n```bash\nMODEL_FLAGS=\"--image_size 64 --num_channels 128 --num_res_blocks 2 --enc_attn_dim 512 \\\n--dropout 0.1 --use_fp16 True\"\nDIFFUSION_FLAGS=\"--noise_schedule cosine\"\nTRAIN_FLAGS=\"--lr 3e-4 --batch_size 512 --microbatch 64\"\n```\nOnce you have setup your hyper-parameters, you can run an experiment like so:\n\n```bash\npython scripts/train_genrator.py --data_dir path/to/images $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS\n```\n\n##### Cascaded training\n* For generator, you also need to modify `small_size` to the size of the input image. \n* Setting `--gaussian_blur True` should be beneficial.\n\n### Classifier\n```bash\nCLASSIFIER_FLAGS=\"--image_size 64 --classifier_depth 2 --classifier_width 128 \\\n--classifier_enc_attn_dim 512 --use_fp16 True\"\nDIFFUSION_FLAGS=\"--noise_schedule cosine\"\nTRAIN_FLAGS=\"--iterations 300000 --batch_size 128 --lr 3e-4 --weight_decay 0.1 \\\n--\"\n```\nAs in generator training, you can run an experiment for classifier as:\n\n```bash\npython scripts/train_classifier.py --data_dir path/to/train_data \\\n--val_data_dir path/to/val_data $CLASSIFIER_FLAGS $TRAIN_FLAGS\n```\n##### Cascaded training\n* For this, it suffices to modify `image_size`, model size, `batch_size`, etc. and run with the exact same command.\n\n##### Remarks\n* We have added [gradient caching](https://github.com/luyug/GradCache) option to enable effective gradient accumulation with contrastive loss.\n* Unlike the generator, overfitting is more likely. Hence, it is recommended to take the checkpoint with the best validation loss.\n\n### Distributed training\n\nYou may also want to train in a distributed manner. In this case, run the same command with `mpiexec`:\n\n```bash\nmpiexec -n $NUM_GPUS python scripts/train_generator.py --data_dir path/to/images \\\n$MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS\n```\nClassifier training can be run likewise.\n\n## Sampling\n\nOnce you have a path to your trained model, you can generate a large batch of samples like so:\n```bash\nMODEL_FLAGS=\"--image_size 64 --num_channels 128 --num_res_blocks 2 --use_fp16 True\"\nDIFFUSION_FLAGS=\"--noise_schedule cosine --timestep_respacing 250\"\nCLASSIFIER_FLAGS=\"--classifier_depth 2 --classifier_width 128 --classifier_enc_attn_dim 512 \\\n--classifier_scale 1.0\"\n```\n```bash\npython scripts/sample.py --model_path /path/to/model_ema.pt \\\n--classifier_path /path/to/classifier.pt --data_dir path/to/texts $MODEL_FLAGS \\\n$CLASSIFIER_FLAGS $DIFFUSION_FLAGS\n```\nYou can remove the relevant parts if classifier-augmented sampling is not used.\n\nAgain, this will save results to a logging directory. Samples are saved as a large `npz` file. A small subset of the samples is also saved as a grid image in jpg and a txt file. \n\nJust like for training, you can run `sample.py` through MPI to use multiple GPUs and machines.\n\nYou can change the number of sampling steps using the `--timestep_respacing` argument. For example, `--timestep_respacing 250` uses 250 steps to sample. Passing `--timestep_respacing ddim25` is similar, but uses the uniform stride from the [DDIM paper](https://arxiv.org/abs/2010.02502).\n\nTo sample using [DDIM](https://arxiv.org/abs/2010.02502), pass `--use_ddim True`.\n\n## Major hyperparameters\n\n* `image_size`: the resolution of the output image.\n* `num_channels`: the number of channels of the outermost layer of U-Net. \n* `enc_attn_dim`: the number of channels of the Transformer encoder.\n* `num_res_blocks`: the number of layers for each resolution of U-Net.\n* `dropout`: recommended to set this to 0.1 for 64 x 64 and 0 otherwise.\n* `noise_schedule`: `cosine` is recommended for `image_size` = 64, and `linear` is recommended otherwise.\n*  `lr`: the base learning rate. The rate is constant for generator and cosine annealed for classifier with linear warmup. \n* `batch_size`: batch size per core.\n* `microbatch`: if set, gradient accumulation is performed with each microbatch size = `microbatch`. Setting this is not recommended for classifier training.\n* `resume_checkpoint`:  path to model parameter checkpoint to resume training (e.g. `--resume_checkpoints path/to/log_dir/model010000_ema.pt`)\n* `text_length`: the length of input text (default 48). The texts longer or shorter than `text_length` are curtailed or padded to this length.  \n\n## Models and Hyperparameters (WIP)\n\nFor model checkpoints (if available) and run flags we have attempted, please refer to [models_hparams.md]() (not ready yet).\n\n## FAQ\n* How can we scale up the model with this repo?\n\t* Model parallelism: \n\t\t* While our repo does not allow model parallelism yet, with cascading training and classifier training, we have multiple models that we can train separately in parallel. This allows a ~4B model without model parallelism on A100s.\n\t* Data parallelism:\n\t\t* Since critical batch size of DDPM seems to be rather large, we can aggressively utilize the data parallelism (e.g. maybe up to `batch_size = 2048` for 64 x 64 generator). \n\t* Dataset size:\n\t\t* Given that even base model requires several hundreds of millions of images (counting multiplicity) for nearly compute-optimal training, and that typical dataset size used for DDPM variants is no more less than 2M, using a dataset of the order of 100M images should improve the performance substantially.\n\t\n## TODO\n- [ ] Finish core components of Diff-DALLE\n\t- [x] Test generator training without classifier\n\t- [x] Test classifier training \n\t- [x] Test sampling without classifier\n\t- [ ] Test sampling with classifier (no bug; performance not checked)\n- [ ] Add the code for preparing a small dataset for demo\n- [ ] Perform large-scale training\n- [ ] Add more details on parameters, compute, dataset size ...\n- [ ] Add evaluation metrics (e.g. Precision \u0026 Recall, FID, etc) and perform ablations\n- [ ] Improve the documentation\n- [ ] Release a pretrained model, colab notebook and web demo \n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FAranKomat%2FDiff-DALLE","html_url":"https://awesome.ecosyste.ms/projects/github.com%2FAranKomat%2FDiff-DALLE","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2FAranKomat%2FDiff-DALLE/lists"}