Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
UNet diffusion model in pure CUDA
| Date | Stars |
|---|---|
| 2026-07-31 | 661 |
| 2026-08-01 | 661 |
| 2026-08-06 | 661 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# unet.cu TL;DR: - UNet diffusion model training written in pure C++/CUDA (only unconditional diffusion right now). - Currently end to end training runs at about 40% the speed of PyTorch with `torch.compile`. The following are benchmarks on one RTX 4090 GPU: | Setup | one full training loop (ms) | | --- | --- | | This repo | 142.44 | | PyTorch | 66.73 | | PyTorch with `torch.compile` | 59.20 | ## Table of contents - [unet.cu](#unetcu) - [Table of contents](#table-of-contents) - [Quick start](#quick-start) - [Introduction](#introduction) - [Background](#background) - [Diffusion models](#diffusion-models) - [UNet architecture](#unet-architecture) - [Version 1: naive implementation](#version-1-naive-implementation) - [Kernels taken from llm.c: linear, groupnorm, and attention](#kernels-taken-from-llmc-linear-groupnorm-and-attention) - [New kernels: upsample, downsample, and convolutions](#new-kernels-upsample-downsample-and-convolutions) - [Benchmarking the first kernel](#benchmarking-the-first-kernel) - [Version 2: custom convolution kernels](#version-2-custom-convolution-kernels) - [`3x3` convolution in detail](#3x3-convolution-in-detail) - [Background on CUDA](#background-on-cuda) - [Kernel details](#kernel-details) - [Profiling version 2](#profiling-version-2) - [Version 3: convolution kernels with vectorized loads](#version-3-convolution-kernels-with-vectorized-loads) - [Future directions](#future-directions) - [Forward pass](#forward-pass) - [Backward pass](#backward-pass) - [Other kernels](#other-kernels) - [Acknowledgements](#acknowledgements) ## Quick start To train a diffusion model in CUDA with some sample images from [ImageNet 64x64](https://image-net.org/index.php), run the following: ```bash gunzip data/elephant_train.bin.gz # prepare the data python train_unet.py --init_model_only True # need to initialize model weights via python make train_unet ./train_unet ``` To train the model with your own data, you need to create a `.bin` file with your data first: ```bash python prepare_data.py --data_dir YOUR_DATA_DIR --output_name YOUR_BINARY_DATA_FILENAME.bin # now run training, assuming you have already initialized the model as above ./train_unet --data_file YOUR_BINARY_DATA_FILENAME.bin ``` The PyTorch training code is essentially taken from the [guided-diffusion](https://github.com/openai/guided-diffusion) repo. To run PyTorch training, do: ```bash python train_unet.py --data_dir YOUR_DATA_DIR # use --compile 0 if you don't want to call torch.compile() on the model ``` The CUDA training loop will save model weights in `.bin` files. To generate new images with model weights saved in either `.bin` or `.pt` files, run: ```bash python generate.py --model_filename YOUR_MODEL_WEIGHTS_FILENAME ``` ## Introduction Inspired by Andrej Karpathy's [llm.c](https://github.com/karpathy/llm.c), I built a UNet from scratch in C/CUDA. The goal of the project is to learn the concepts in llm.c, and to reach for PyTorch's performance with our CUDA implementation. I chose the UNet because it is a key architecture for diffusion models, and I will do some simple diffusion model training with it. Diffusion model training is quite sophisticated nowadays. Since this project is focused on learning CUDA as opposed to building the best diffusion model, I prioritized simplicity over performance, and followed the implementation from the paper [Diffusion Models Beat GANs on Image Synthesis](https://arxiv.org/abs/2105.05233). Currently the UNet only supports unconditioned diffusion training. I also did not reproduce all the model configurations from the paper; the details of the differences will be explained in the [section](#unet-architecture) on the architecture. Here are some images generated with our CUDA implementation. The model is trained on elephant images from [ImageNet 64x64](https://image-net.org/index.php) without class-conditioning. The model is highly over fitting the training set right
Excerpt of 39,450 characters
Read on GitHub
Chen Lu
7
Would you bet a product on this? Bounded 0–100 and slow moving.
matched fp:869fe0cba99a352d, desc:diffusion model