adding dit

Author: codingwithsurya

README.md

@@ -1,16 +1,16 @@
 # diffusion.cu
 
-This project is a from-scratch implementation of diffusion model training in C++/CUDA. This project is currently in-progress. Inspired by Andrej Karpathy's [llm.c](https://github.com/karpathy/llm.c) and Chen Lu's [unet.cu](https://github.com/clu0/unet.cu). The implementation is based on the architecture in the paper [Diffusion Models Beat GANs on Image Synthesis](https://arxiv.org/abs/2105.05233).
+This project is a from-scratch implementation of diffusion model training in C++/CUDA. This project is currently in-progress, and supports the classic UNet architecture, based on the architecture in the paper [Diffusion Models Beat GANs on Image Synthesis](https://arxiv.org/abs/2105.05233), and the transformer architecture (DiT), based on [Scalable Diffusion Models with Transformers] (https://arxiv.org/abs/2212.09748). Inspired by Andrej Karpathy's [llm.c](https://github.com/karpathy/llm.c) and Chen Lu's [unet.cu](https://github.com/clu0/unet.cu).
 
 ## Training
 
-You can train it on the images from the ImageNet 64x64 dataset via 
+UNet currently supports training. You can train it on the images from the ImageNet 64x64 dataset via 
 
 ```
-gunzip data/elephant_train.bin.gz 
-python train_diffusion.py --init_model_only True 
-make train_diffusion
-./train_diffusion`
+gunzip unet/data/elephant_train.bin.gz 
+python unet/train_diffusion.py --init_model_only True 
+make -C unet train_diffusion
+./unet/train_diffusion
 ```
 
 ### **Current Implementation:**
@@ -26,7 +26,6 @@ This currently supports unconditional diffusion model training, and the end-to-e
 In Progress:
 - support for distributed training via MPI
 - support for mixed precision training
-- support for DiT as another architecture we can add, because we can re-use the same components of llm.c
 
 ### **My Motivation:**