This is the repository for ZJU 课程综合实践Ⅱ 2023, the project we select is Image generation based on Diffusion Module, the implementation is based on DDPM and DDIM model, we use Kaggle Dataset to train our implementation of UNet, and it has achieved better results for generating images.
We also simplify UNet by removing Attention Block , Middle Block and reducing the number of convolutional layers to train single image, it achieves interesting generation results different from UNet.
- DDPM:
![](imgs/ddpm.png)
- DDIM:
- Single Diffusion:
- model:
model
├── ddim.py
├── ddpm.py
├── sunet.py
└── unet.py
ddpm.py
: the class for ddpm model, including processes such as denoising and adding noise and some other helper functions.
ddim.py
: Inherited from ddpm, implement the denoising process of ddim.
unet.py
: the implementation of unet, We adapted our model based on the original unet and made some adjustments.
sunet.py
: a simplified implementation for unet, it doesn't use middle block、attention block and reduced neural network depth, it will be trained by Single Diffusion Model.
- dataset:
dataset
├── kaggle.py
└── single.py
kaggle.py
: dataset for Anime Names and Images Dataset.
single.py
: dataset for Single Diffusion Model.
-
Generated images of NetWork trained by Anime Names and Images Dataset
-
Our sunet trained by Single Image
We train our model/sunet.py
by single image:
There are results for Generated images, you can see even training with a single image produces images with different effects
You can see the Genrate steps for images:
- Linux
- Python3
- CPU or NVIDIA GPU
We use Anime Names and Images Dataset by Kaggle to train our naive DDPM/DDIM model.
We use single image online to train our Single Diffusion Model, you can see the image in data/single/*.png
.
Actually, you can use the datasets/images to train our models.
- Training the model:
To train the model for Kaggle Dataset, execute the following command.
Since login permission is required to download the dataset, if you want to train
Kaggle Dataset, you should download it first, then named ./arhive/dataset/dataset/*.jpg
as ./kaggle/*.jpg
to replace the folder ./data/kaggle
and use sh get_txt.sh
to get a new name.txt
.
If you want to use the pre-trained network, you can use --load
,else you can
train a new network yourself.
python3 train_kaggle.py\
--dataset ./data/ \
--img_dir kaggle/ \
--data_txt name.txt \
--epoch=32 \
--image_size=64 \
--b=16 \
--load \
--model_name Gen8000.pth
To train the Single Diffusion model , execute the following command.
python3 train_single.py \
--dataset ./data/single/ \
--img_name geese.png \
--epoch=32 \
--image_size=128 \
--b=32 \
--load \
--model_name geese.pth
- Generate the images:
To generate images trained by Kaggle DataSet, execute the following command.
The images will be generated by ddpm and ddim, you can set --ddim_steps
for the steps of ddim.
python3 gen_kaggle.py\
--dataset ./data/ \
--img_dir kaggle/ \
--data_txt name.txt \
--image_size=64 \
--load \
--model_name Gen8000.pth \
--gen_name gen \
--ddim_steps=100
To generate images by Single Diffusion model , execute the following command.
python3 gen_single.py \
--dataset ./data/single/ \
--img_name geese.png \
--image_size=128 \
--load \
--model_name geese.pth
We supply our pretrained NetWork in ..., you can use these Networks to generate the images
Our Neural Networks are based on UNet
1.DDPM:🔗
2.DDIM:🔗
3.Single Diffusion:🔗
4.DDPM PyTorch implementation:🔗