Awesome

SimpleSDM

This repository contains a simple and flexible PyTorch implementation of StableDiffusion based on diffusers.

<div align="center"> <img src="example/example.png"> </div>

Prepartion

You should download the checkpoints of SDM-1.5, from SDM-1.5, including scheduler, text_encoder, tokenizer, unet, and vae. Then put it in the ckpt folder.

Requirements

Python >= 3.8 (Recommend to use Anaconda or Miniconda)
PyTorch >= 1.12
xformers == 0.0.13
diffusers == 0.13.1
accelerate == 0.17.1
transformers == 4.27.4

A suitable conda environment named ldm can be created and activated with:

conda env create -f environment.yaml
conda activate ldm

Dataset Preparation

You need write a DataLoader suitable for your own Dataset, because we just provide a simple example to test the code.

Training

CUDA_VISIBLE_DEVICES=0,1,2 accelerate launch --multi_gpu train.py

Inference

CUDA_VISIBLE_DEVICES=0 python inference.py --prompt "A cat is running in the rain."

Acknowledgements

Many thanks to the code bases from diffusers.