From d03f4429c1155eb806d9b0763a43dfe4184a98f9 Mon Sep 17 00:00:00 2001
From: Fazzie-Maqianli <55798671+Fazziekey@users.noreply.github.com>
Date: Wed, 15 Feb 2023 09:55:53 +0800
Subject: [PATCH] add ci (#2641)

---
 examples/images/diffusion/README.md  |  6 ++++++
 examples/images/diffusion/main.py    |  2 ++
 examples/images/diffusion/test_ci.sh | 17 +++++++++++++++++
 3 files changed, 25 insertions(+)
 mode change 100644 => 100755 examples/images/diffusion/test_ci.sh

diff --git a/examples/images/diffusion/README.md b/examples/images/diffusion/README.md
index 952da5d1c..15932f1f5 100644
--- a/examples/images/diffusion/README.md
+++ b/examples/images/diffusion/README.md
@@ -92,6 +92,12 @@ cd ColossalAI
 CUDA_EXT=1 pip install .
 ```
 
+#### Step 3:Accelerate with flash attention by xformers(Optional)
+
+```
+pip install xformers
+```
+
 ### Option #2: Use Docker
 
 To use the stable diffusion Docker image, you can either build using the provided the [Dockerfile](./docker/Dockerfile) or pull a Docker image from our Docker hub.
diff --git a/examples/images/diffusion/main.py b/examples/images/diffusion/main.py
index 5f166aa1f..4dd88a5ec 100644
--- a/examples/images/diffusion/main.py
+++ b/examples/images/diffusion/main.py
@@ -539,6 +539,8 @@ if __name__ == "__main__":
         raise ValueError("-n/--name and -r/--resume cannot be specified both."
                          "If you want to resume training in a new log folder, "
                          "use -n/--name in combination with --resume_from_checkpoint")
+
+    ckpt = None
     if opt.resume:
         rank_zero_info("Resuming from {}".format(opt.resume))
         if not os.path.exists(opt.resume):
diff --git a/examples/images/diffusion/test_ci.sh b/examples/images/diffusion/test_ci.sh
old mode 100644
new mode 100755
index e69de29bb..51ceeb41d
--- a/examples/images/diffusion/test_ci.sh
+++ b/examples/images/diffusion/test_ci.sh
@@ -0,0 +1,17 @@
+#!/bin/bash
+set -euxo pipefail
+
+conda env create -f environment.yaml
+
+conda activate ldm
+
+conda install pytorch==1.12.1 torchvision==0.13.1 torchaudio==0.12.1 cudatoolkit=11.3 -c pytorch
+pip install transformers diffusers invisible-watermark
+
+CUDA_EXT=1  pip install colossalai
+
+pip install pytorch-lightning
+
+wget https://huggingface.co/stabilityai/stable-diffusion-2-base/resolve/main/512-base-ema.ckpt
+
+python main.py --logdir /tmp --train --base configs/Teyvat/train_colossalai_teyvat.yaml --ckpt 512-base-ema.ckpt