{"id":14958116,"url":"https://github.com/vachanvy/diffusion-transformer","last_synced_at":"2026-03-10T16:03:35.936Z","repository":{"id":253667819,"uuid":"843285544","full_name":"VachanVY/diffusion-transformer","owner":"VachanVY","description":"Pytorch and JAX Implementation of Scalable Diffusion Models with Transformers | Diffusion Transformers in Pytorch and JAX","archived":false,"fork":false,"pushed_at":"2024-10-16T05:43:16.000Z","size":59240,"stargazers_count":4,"open_issues_count":0,"forks_count":0,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-01-31T02:02:37.025Z","etag":null,"topics":["deep-learning","diffusion-models","diffusion-transformer","jax","jax-implementation","python","pytorch","pytorch-implementation","transformers"],"latest_commit_sha":null,"homepage":"","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/VachanVY.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,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-08-16T07:16:40.000Z","updated_at":"2024-12-28T18:23:58.000Z","dependencies_parsed_at":"2024-08-29T13:40:17.786Z","dependency_job_id":null,"html_url":"https://github.com/VachanVY/diffusion-transformer","commit_stats":null,"previous_names":["vachanvy/diffusion-transformer"],"tags_count":0,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/VachanVY%2Fdiffusion-transformer","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/VachanVY%2Fdiffusion-transformer/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/VachanVY%2Fdiffusion-transformer/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/VachanVY%2Fdiffusion-transformer/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/VachanVY","download_url":"https://codeload.github.com/VachanVY/diffusion-transformer/tar.gz/refs/heads/main","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":237982289,"owners_count":19397236,"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":["deep-learning","diffusion-models","diffusion-transformer","jax","jax-implementation","python","pytorch","pytorch-implementation","transformers"],"created_at":"2024-09-24T13:16:15.078Z","updated_at":"2026-03-10T16:03:35.923Z","avatar_url":"https://github.com/VachanVY.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# Diffusion-Transformers\n## Contents\n* [**CelebA**](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#celeba) \n   * **[Generated-images](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#generated-images)** \u003c====== See the Model Generated Images here\n   * **[Training-insights](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#training-insights)**\n* **[MNIST-experiment](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#mnist-experiment)**\n   * [**Training on MNIST**](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#training-on-mnist)\n* **[Diffusion-Transformers Paper Summary](https://github.com/VachanVY/diffusion-transformer?tab=readme-ov-file#latent-diffusion-models)**\n\n---\n* To clone the repo and install libraries from requirements.txt run the below commands\n   ```python\n   git clone https://github.com/VachanVY/diffusion-transformer.git\n   pip install -r requirements.txt\n   ```\n---\n## CelebA\n* \u003cimg src=\"images/_____image.png\" alt=\"CelebA Sample\" width=\"400\"/\u003e\\\n   The CelebA dataset consists of celebrity images, such as the example shown above.\n* *The model has been trained for only 100K steps so far. Ideally, it should be trained for 400K steps to improve the quality of generated images. While the current model is undertrained, the generated images are still decent*. Check them out below.\n* ***If you have GPUs and can train this model for 400k steps, please edit the generate.py file to include a download link to your weights, and send a pull request. I’d be happy to incorporate it! You can start training from the checkpoint, download it using the below command***\n  ```python\n   python checkpoints/celeba/download_celeba_ckpt.py\n  ```\n\n### Generated Images\n* Run the following command to generate images:\n   ```python\n   python checkpoints/celeba/download_celeba_ckpt.py # download weights\n   python generate.py --labels female male female female --ema True # generate images\n   ```\n* Some generated images:\\\n    \u003cimg src=\"https://github.com/user-attachments/assets/6cbe6bc7-1e7a-44ed-83ae-df64c40bf8d4\" alt=\"Alt text\" width=\"300\"\u003e\n    \u003cimg src=\"https://github.com/user-attachments/assets/181f31fe-4b4c-4719-93df-f3d767853608\" alt=\"Alt text\" width=\"300\"\u003e\n    \u003cimg src=\"https://github.com/user-attachments/assets/eb463c13-45b2-474e-a431-b3b296fe00c4\" alt=\"Alt text\" width=\"300\"\u003e\n    \u003cimg src=\"https://github.com/user-attachments/assets/2048e0ba-4ebd-48f1-a379-df39680495f9\" alt=\"Alt text\" width=\"300\"\u003e\n    \u003cimg src=\"https://github.com/user-attachments/assets/e7bea3b5-ce52-4de6-a03e-a7694e66b320\" alt=\"Alt text\" width=\"300\"\u003e\n    \u003cimg src=\"https://github.com/user-attachments/assets/fe7d1f7d-f7a8-4059-a1f2-172a0e2345d5\" alt=\"Alt text\" width=\"300\"\u003e\n\n\n### Training Insights\n* ![www](images/loss_vs_steps.png)\n* Run the following file to get the graph of `loss` vs `learning_rate` to select the best `max_lr`:\n  ```python\n  python torch_src/config.py\n  ```\n  ![d34r3r](images/loss_vs_lr.png)\n  You can see that \"log10 learning rates\" after `-4.0 (lr=1e-4)` are unstable. When using a max_lr=3e-4, the training loss spiked, and the model forgot everything midway through training\n  ```python\n  python train.py\n  ```\n* A GIF showing images as the training progress is displayed below\n ![Animated GIF](images/training_progress.gif)\n\n\n## MNIST Experiment\n* The MNIST dataset is used as a test case for the diffusion model.\n\n\n### Training on MNIST\n* To train the model on MNIST, run:\n   ```python\n   python torch_src/diffusion_transformer.py\n   ```\n* Additionally, check out the [jax_mnist_diffusion_transformer.ipynb](https://nbviewer.org/github/VachanVY/diffusion-transformer/blob/main/jax_mnist_diffussion_transformer.ipynb) notebook for a JAX version of the model:\n![alt text](images/______image-1.png)\n\n\n## Latent-Diffusion Models\n* In this paper, we apply DiTs to latent space, although they could be applied to pixel space without modification as well\n\n## Classifier-free Guidance\n![alt text](images/dit-1.png)\n\n## DiTs\n* ![alt text](images/dit.png)\n* Following patchify, we apply standard ViT frequency-based positional embeddings (the sine-cosine version) to all input tokens\n![alt text](images/dit-2.png)\n\n### DiT block design\nWe explore four variants of transformer blocks that process conditional inputs differently\n* ***In-context conditioning:*** We simply append the vector\nembeddings of t and c as two additional tokens in\nthe input sequence, treating them no differently from\nthe image tokens. This is similar to cls tokens in\nViTs, and it allows us to use standard ViT blocks without\nmodification.\n* ***Cross-attention block:*** We concatenate the embeddings\nof t and c into a length-two sequence, separate from\nthe image token sequence. The transformer block is\nmodified to include an additional multi-head crossattention\nlayer following the multi-head self-attention\nblock, similar to the original design from Vaswani et\nal., and also similar to the one used by LDM for\nconditioning on class labels. Cross-attention adds the\nmost Gflops to the model, roughly a 15% overhead.\n* ***Adaptive layer norm (adaLN) block:*** We explore replacing standard layer norm layers in transformer blocks with adaptive layer norm (adaLN). Rather than directly learn dimensionwise scale and shift parameters $\\gamma$ and $\\beta$, we regress\nthem from the sum of the embedding vectors of $t$ and\n$c$. Of the three block designs we explore, adaLN adds\nthe least Gflops and is thus the most compute-efficient\n* ***adaLN-Zero block*** : Zero-initializing the final batch norm scale factor in each block accelerates large-scale training in the supervised learning setting. Diffusion U-Net models use a similar initialization strategy, zero-initializing the final convolutional layer in each block prior to any residual connections. We explore a modification of the adaLN DiT block which does the same. In addition to regressing $\\gamma$ and $\\beta$, we also regress dimensionwise scaling parameters $\\alpha$ that are applied immediately prior to any residual connections within the DiT block\n\n## Training Setup\n* We initialize the final linear layer with zeros and otherwise use standard weight initialization techniques from ViT\n* AdamW\n* We use a constant learning rate of $1\\times 10^{-4}$, no weight decay and a batch size of $256$\n* Exponential moving average (EMA) of DiT weights over training with a decay of $0.9999$\n* The VAE encoder has a downsample factor of 8, given an\nRGB image $x$ with shape $256 \\times 256 \\times 3$, $z = E(x)$ has shape $32 \\times 32 \\times 4$\n* $t_{\\max} = 1000$ linear variance schedule ranging from $1 \\times 10^{-4}$ to $2 \\times 10^{-2}$\n\n# Diffusion Models—Theory\n## Variational Autoencoder (VAE) Loss Derivation\n* ![WhatsApp Image 2025-06-19 at 21 29 16_3537a13f](https://github.com/user-attachments/assets/ae1d56df-8390-4915-bdcf-b94ad6bfc380)\n\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fvachanvy%2Fdiffusion-transformer","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fvachanvy%2Fdiffusion-transformer","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fvachanvy%2Fdiffusion-transformer/lists"}