haoningwu3639/SimpleSDM

A simple and flexible PyTorch implementation of StableDiffusion based on diffusers.

SimpleSDM

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

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.