Keras Stable Diffusion: GPU starter example

GitHub：https://github.com/divamgupta/stable-diffusion-tensorflow

Install GPU requirements

pip install git+https://github.com/fchollet/stable-diffusion-tensorflow --upgrade --quiet
pip install tensorflow tensorflow_addons ftfy --upgrade --quiet
apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2

Let's instantiate a Text2Image generator and make a first image

The first run has a bit of extra compilation overhead.

from stable_diffusion_tf.stable_diffusion import Text2Image
from PIL import Image
 
generator = Text2Image( 
    img_height=512,
    img_width=512,
    jit_compile=False,  # You can try True as well (different performance profile)
)
img = generator.generate(
    "DSLR photograph of an astronaut riding a horse",
    num_steps=50,
    unconditional_guidance_scale=7.5,
    temperature=1,
    batch_size=1,
)
pil_img = Image.fromarray(img[0])
display(pil_img)

We can keep making more images with the same generator

You only have to go through compilation once -- all subsequent runs are faster.

img = generator.generate(
    "An epic unicorn riding in the sunset, artstation concept art",
    num_steps=50,
    unconditional_guidance_scale=7.5,
    temperature=1,
)
pil_img = Image.fromarray(img[0])
display(pil_img)

Let's try batched generation

img = generator.generate(
    "Ruins of a castle in Scotland",
    num_steps=50,
    unconditional_guidance_scale=7.5,
    temperature=1,
    batch_size=4,
)
pil_img = Image.fromarray(img[0])
display(pil_img)

pil_img = Image.fromarray(img[1])
display(pil_img)

pil_img = Image.fromarray(img[2])
display(pil_img)

在 WSL2 中食用方法（主要是 GPU 加速）

除上面提到了的几个包之外还需要执行：

pip install tensorflow-directml-plugin

在 WSL 中执行：

nvidia-smi

结果：

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.75       Driver Version: 517.40       CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0  On |                  N/A |
| N/A   59C    P8    20W /  N/A |   2312MiB /  8192MiB |      8%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
 
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

按照自己的 WSL 版本和 Cuda 版本到 https://developer.nvidia.com/rdp/cudnn-archive 下载相应的驱动

下载后安装：

sudo dpkg -i cudnn-local-repo-ubuntu2004-8.4.1.50_1.0-1_amd64.deb

如果报错：

The public CUDA GPG key does not appear to be installed.
To install the key, run this command:
sudo cp /var/cudnn-local-repo-ubuntu2004-8.4.1.50/cudnn-local-E3EC4A60-keyring.gpg /usr/share/keyrings/

执行：

sudo cp /var/cudnn-local-repo-ubuntu2004-8.4.1.50/cudnn-local-E3EC4A60-keyring.gpg /usr/share/keyrings/cuda-archive-keyring.gpg

如果是较早版本的 Cuda 可能会报错：

/sbin/ldconfig.real: /usr/lib/wsl/lib/libcuda.so.1 is not a symbolic link

解决方法：将以下内容写到 /etc/wsl.conf 文件中

[automount]
ldconfig = false

tips

可能还需要安装 tqdm, regex, pillow 几个包