基于 huggingface diffuser 库本地部署 Stable diffusion

前言

Stable Diffusion 是用 LAION-5B 的子集(图像大小为512*512)训练的扩散模型。此模型冻结 CLIP 的 ViT-L/14 文本编码器建模 prompt text。模型包含 860M UNet 和123M 文本编码器,可运行在具有至少10GB VRAM 的 GPU 上。

HF主页:https://huggingface.co/CompVis/stable-diffusion
Colab:https://colab.research.google.com/github/huggingface/notebooks/blob/main/diffusers/stable_diffusion.ipynb
diffusers官方文档:https://huggingface.co/docs/diffusers

接下来实战一下本地部署。

本地部署

1. 安装环境

conda create -n diffenv python=3.8
conda activate diffenv
pip install diffusers==0.4.0
pip install transformers scipy ftfy
# pip install "ipywidgets>=7,<8" 这个是colab用于交互输入的控件

如果后面执行代码时报错 RuntimeError: CUDA error: no kernel image is available for execution on the device,说明cuda版本和pytorch版本问题,根据机器的 cuda 版本重新装一下:

pip install torch==1.12.0+cu113 torchvision==0.13.0+cu113 torchaudio==0.12.0 -f https://download.pytorch.org/whl/torch_stable.html

获取模型:
首先得同意模型的使用协议。
如果用官方 colab,需要输入 huggingface 的 access token 来联网校验你是否同意了协议。如果不想输入的话,就执行以下命令先把模型权重等文件下载到本地:

git lfs install
git clone https://huggingface.co/CompVis/stable-diffusion-v1-4

这样加载模型时直接 DiffusionPipeline.from_pretrained("./MODEL_PATH/stable-diffusion-v1-4") ,就不用加 use_auth_token=AUTH_TOKEN 参数了。

2. 加载模型

如果要确保高精度(占显存也高),删除 revision="fp16"torch_dtype=torch.float16

import torch, os
from diffusers import StableDiffusionPipeline

os.environ["CUDA_VISIBLE_DEVICES"] = "2"

pipe = StableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",  # 本地地址也行
    revision="fp16",  # 如果不想用半精度,删掉这行和下面一行
    torch_dtype=torch.float16
)  

pipe = pipe.to("cuda")

3. 生成图像

3.1 直接生成

默认长宽都是512像素,可以指定 pipe(height=512, width=768) 来控制尺寸。需要注意的是:

  • 高度和宽度都是8的倍数
  • 低于 512 可能会导致图像质量降低
  • 长宽都超过 512 会出现重复图像区域(丢失全局一致性)
  • 非正方形图像最佳方法:一个维度为 512px,另一维度大于512px。
prompt = "a photograph of an astronaut swimming in the river"
image = pipe(prompt).images[0]  # PIL格式 (https://pillow.readthedocs.io/en/stable/)
image.save(f"astronaut_rides_horse.png")
image

输出结果如下,还有点内味儿。
基于 huggingface diffuser 库本地部署 Stable diffusion_第1张图片

3.2 非随机生成

刚才 3.1 部分生成的每次都不一样,若需非随机生成,则指定随机种子,pipe() 中传入 generator 参数指定 generator。

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, generator=generator).images[0]
image

基于 huggingface diffuser 库本地部署 Stable diffusion_第2张图片

3.3 推理步数控制图像质量

使用 num_inference_steps 参数更改推理 steps。通常步数越多,结果越好,推理越慢。Stable Diffusion 比较强,只需相对较少的步骤效果就不错,因此建议使用默认值50。如图把 num_inference_steps 设成 100,随机种子保持不变,貌似效果差距并不大。

import torch

generator = torch.Generator("cuda").manual_seed(1024)
image = pipe(prompt, num_inference_steps=100, generator=generator).images[0]
image

基于 huggingface diffuser 库本地部署 Stable diffusion_第3张图片

3.4 生成多张图片

写个做图片拼接的函数:

from PIL import Image

def image_grid(imgs, rows, cols):
    assert len(imgs) == rows*cols

    w, h = imgs[0].size
    grid = Image.new('RGB', size=(cols*w, rows*h))
    grid_w, grid_h = grid.size
    
    for i, img in enumerate(imgs):
        grid.paste(img, box=(i%cols*w, i//cols*h))
    return grid

一次性生成 3 幅图,此时 prompt 为 list 而不是 str。

num_images = 3
prompt = ["a traditional Chinese painting of a squirrel eating a banana"] * num_images
images = pipe(prompt).images
grid = image_grid(images, rows=1, cols=3)
grid

基于 huggingface diffuser 库本地部署 Stable diffusion_第4张图片

附录

依赖环境如下:

Package                       Version
----------------------------- ----------------------------
absl-py                       1.2.0
aeppl                         0.0.33
aesara                        2.7.9
aiohttp                       3.8.3
aiosignal                     1.2.0
alabaster                     0.7.12
albumentations                1.2.1
altair                        4.2.0
appdirs                       1.4.4
arviz                         0.12.1
astor                         0.8.1
astropy                       4.3.1
astunparse                    1.6.3
async-timeout                 4.0.2
asynctest                     0.13.0
atari-py                      0.2.9
atomicwrites                  1.4.1
attrs                         22.1.0
audioread                     3.0.0
autograd                      1.5
Babel                         2.10.3
backcall                      0.2.0
beautifulsoup4                4.6.3
bleach                        5.0.1
blis                          0.7.8
bokeh                         2.3.3
branca                        0.5.0
bs4                           0.0.1
CacheControl                  0.12.11
cached-property               1.5.2
cachetools                    4.2.4
catalogue                     2.0.8
certifi                       2022.9.24
cffi                          1.15.1
cftime                        1.6.2
chardet                       3.0.4
charset-normalizer            2.1.1
click                         7.1.2
clikit                        0.6.2
cloudpickle                   1.5.0
cmake                         3.22.6
cmdstanpy                     1.0.7
colorcet                      3.0.1
colorlover                    0.3.0
community                     1.0.0b1
confection                    0.0.2
cons                          0.4.5
contextlib2                   0.5.5
convertdate                   2.4.0
crashtest                     0.3.1
crcmod                        1.7
cufflinks                     0.17.3
cupy-cuda11x                  11.0.0
cvxopt                        1.3.0
cvxpy                         1.2.1
cycler                        0.11.0
cymem                         2.0.6
Cython                        0.29.32
daft                          0.0.4
dask                          2022.2.0
datascience                   0.17.5
debugpy                       1.0.0
decorator                     4.4.2
defusedxml                    0.7.1
descartes                     1.1.0
diffusers                     0.4.0
dill                          0.3.5.1
distributed                   2022.2.0
dlib                          19.24.0
dm-tree                       0.1.7
docutils                      0.17.1
dopamine-rl                   1.0.5
earthengine-api               0.1.326
easydict                      1.10
ecos                          2.0.10
editdistance                  0.5.3
en-core-web-sm                3.4.0
entrypoints                   0.4
ephem                         4.1.3
et-xmlfile                    1.1.0
etils                         0.8.0
etuples                       0.3.8
fa2                           0.3.5
fastai                        2.7.9
fastcore                      1.5.27
fastdownload                  0.0.7
fastdtw                       0.3.4
fastjsonschema                2.16.2
fastprogress                  1.0.3
fastrlock                     0.8
feather-format                0.4.1
filelock                      3.8.0
firebase-admin                4.4.0
fix-yahoo-finance             0.0.22
Flask                         1.1.4
flatbuffers                   22.9.24
folium                        0.12.1.post1
frozenlist                    1.3.1
fsspec                        2022.8.2
ftfy                          6.1.1
future                        0.16.0
gast                          0.5.3
GDAL                          2.2.2
gdown                         4.4.0
gensim                        3.6.0
geographiclib                 1.52
geopy                         1.17.0
gin-config                    0.5.0
glob2                         0.7
google                        2.0.3
google-api-core               1.31.6
google-api-python-client      1.12.11
google-auth                   1.35.0
google-auth-httplib2          0.0.4
google-auth-oauthlib          0.4.6
google-cloud-bigquery         1.21.0
google-cloud-bigquery-storage 1.1.2
google-cloud-core             1.0.3
google-cloud-datastore        1.8.0
google-cloud-firestore        1.7.0
google-cloud-language         1.2.0
google-cloud-storage          1.18.1
google-cloud-translate        1.5.0
google-colab                  1.0.0
google-pasta                  0.2.0
google-resumable-media        0.4.1
googleapis-common-protos      1.56.4
googledrivedownloader         0.4
graphviz                      0.10.1
greenlet                      1.1.3
grpcio                        1.49.1
gspread                       3.4.2
gspread-dataframe             3.0.8
gym                           0.25.2
gym-notices                   0.0.8
h5py                          3.1.0
HeapDict                      1.0.1
hijri-converter               2.2.4
holidays                      0.16
holoviews                     1.14.9
html5lib                      1.0.1
httpimport                    0.5.18
httplib2                      0.17.4
httplib2shim                  0.0.3
httpstan                      4.6.1
huggingface-hub               0.10.0
humanize                      0.5.1
hyperopt                      0.1.2
idna                          2.10
imageio                       2.9.0
imagesize                     1.4.1
imbalanced-learn              0.8.1
imblearn                      0.0
imgaug                        0.4.0
importlib-metadata            5.0.0
importlib-resources           5.9.0
imutils                       0.5.4
inflect                       2.1.0
intel-openmp                  2022.2.0
intervaltree                  2.1.0
ipykernel                     5.3.4
ipython                       7.9.0
ipython-genutils              0.2.0
ipython-sql                   0.3.9
ipywidgets                    7.7.1
itsdangerous                  1.1.0
jax                           0.3.21
jaxlib                        0.3.20+cuda11.cudnn805
jedi                          0.18.1
jieba                         0.42.1
Jinja2                        2.11.3
joblib                        1.2.0
jpeg4py                       0.1.4
jsonschema                    4.3.3
jupyter-client                6.1.12
jupyter-console               6.1.0
jupyter-core                  4.11.1
jupyterlab-widgets            3.0.3
kaggle                        1.5.12
kapre                         0.3.7
keras                         2.8.0
Keras-Preprocessing           1.1.2
keras-vis                     0.4.1
kiwisolver                    1.4.4
korean-lunar-calendar         0.3.1
langcodes                     3.3.0
libclang                      14.0.6
librosa                       0.8.1
lightgbm                      2.2.3
llvmlite                      0.39.1
lmdb                          0.99
locket                        1.0.0
logical-unification           0.4.5
LunarCalendar                 0.0.9
lxml                          4.9.1
Markdown                      3.4.1
MarkupSafe                    2.0.1
marshmallow                   3.18.0
matplotlib                    3.2.2
matplotlib-venn               0.11.7
miniKanren                    1.0.3
missingno                     0.5.1
mistune                       0.8.4
mizani                        0.7.3
mkl                           2019.0
mlxtend                       0.14.0
more-itertools                8.14.0
moviepy                       0.2.3.5
mpmath                        1.2.1
msgpack                       1.0.4
multidict                     6.0.2
multipledispatch              0.6.0
multitasking                  0.0.11
murmurhash                    1.0.8
music21                       5.5.0
natsort                       5.5.0
nbconvert                     5.6.1
nbformat                      5.6.1
netCDF4                       1.6.1
networkx                      2.6.3
nibabel                       3.0.2
nltk                          3.7
notebook                      5.3.1
numba                         0.56.2
numexpr                       2.8.3
numpy                         1.21.6
oauth2client                  4.1.3
oauthlib                      3.2.1
okgrade                       0.4.3
opencv-contrib-python         4.6.0.66
opencv-python                 4.6.0.66
opencv-python-headless        4.6.0.66
openpyxl                      3.0.10
opt-einsum                    3.3.0
osqp                          0.6.2.post0
packaging                     21.3
palettable                    3.3.0
pandas                        1.3.5
pandas-datareader             0.9.0
pandas-gbq                    0.13.3
pandas-profiling              1.4.1
pandocfilters                 1.5.0
panel                         0.12.1
param                         1.12.2
parso                         0.8.3
partd                         1.3.0
pastel                        0.2.1
pathlib                       1.0.1
pathy                         0.6.2
patsy                         0.5.2
pep517                        0.13.0
pexpect                       4.8.0
pickleshare                   0.7.5
Pillow                        7.1.2
pip                           21.1.3
pip-tools                     6.2.0
plotly                        5.5.0
plotnine                      0.8.0
pluggy                        0.7.1
pooch                         1.6.0
portpicker                    1.3.9
prefetch-generator            1.0.1
preshed                       3.0.7
prettytable                   3.4.1
progressbar2                  3.38.0
promise                       2.3
prompt-toolkit                2.0.10
prophet                       1.1.1
protobuf                      3.17.3
psutil                        5.4.8
psycopg2                      2.9.3
ptyprocess                    0.7.0
py                            1.11.0
pyarrow                       6.0.1
pyasn1                        0.4.8
pyasn1-modules                0.2.8
pycocotools                   2.0.5
pycparser                     2.21
pyct                          0.4.8
pydantic                      1.9.2
pydata-google-auth            1.4.0
pydot                         1.3.0
pydot-ng                      2.0.0
pydotplus                     2.0.2
PyDrive                       1.3.1
pyemd                         0.5.1
pyerfa                        2.0.0.1
Pygments                      2.6.1
pygobject                     3.26.1
pylev                         1.4.0
pymc                          4.1.4
PyMeeus                       0.5.11
pymongo                       4.2.0
pymystem3                     0.2.0
PyOpenGL                      3.1.6
pyparsing                     3.0.9
pyrsistent                    0.18.1
pysimdjson                    3.2.0
pysndfile                     1.3.8
PySocks                       1.7.1
pystan                        3.3.0
pytest                        3.6.4
python-apt                    0.0.0
python-chess                  0.23.11
python-dateutil               2.8.2
python-louvain                0.16
python-slugify                6.1.2
python-utils                  3.3.3
pytz                          2022.4
pyviz-comms                   2.2.1
PyWavelets                    1.3.0
PyYAML                        6.0
pyzmq                         23.2.1
qdldl                         0.1.5.post2
qudida                        0.0.4
regex                         2022.6.2
requests                      2.23.0
requests-oauthlib             1.3.1
resampy                       0.4.2
rpy2                          3.4.5
rsa                           4.9
scikit-image                  0.18.3
scikit-learn                  1.0.2
scipy                         1.7.3
screen-resolution-extra       0.0.0
scs                           3.2.0
seaborn                       0.11.2
Send2Trash                    1.8.0
setuptools                    57.4.0
setuptools-git                1.2
Shapely                       1.8.4
six                           1.15.0
sklearn-pandas                1.8.0
smart-open                    5.2.1
snowballstemmer               2.2.0
sortedcontainers              2.4.0
soundfile                     0.11.0
spacy                         3.4.1
spacy-legacy                  3.0.10
spacy-loggers                 1.0.3
Sphinx                        1.8.6
sphinxcontrib-serializinghtml 1.1.5
sphinxcontrib-websupport      1.2.4
SQLAlchemy                    1.4.41
sqlparse                      0.4.3
srsly                         2.4.4
statsmodels                   0.12.2
sympy                         1.7.1
tables                        3.7.0
tabulate                      0.8.10
tblib                         1.7.0
tenacity                      8.1.0
tensorboard                   2.8.0
tensorboard-data-server       0.6.1
tensorboard-plugin-wit        1.8.1
tensorflow                    2.8.2+zzzcolab20220929150707
tensorflow-datasets           4.6.0
tensorflow-estimator          2.8.0
tensorflow-gcs-config         2.8.0
tensorflow-hub                0.12.0
tensorflow-io-gcs-filesystem  0.27.0
tensorflow-metadata           1.10.0
tensorflow-probability        0.16.0
termcolor                     2.0.1
terminado                     0.13.3
testpath                      0.6.0
text-unidecode                1.3
textblob                      0.15.3
thinc                         8.1.2
threadpoolctl                 3.1.0
tifffile                      2021.11.2
tokenizers                    0.12.1
toml                          0.10.2
tomli                         2.0.1
toolz                         0.12.0
torch                         1.12.1+cu113
torchaudio                    0.12.1+cu113
torchsummary                  1.5.1
torchtext                     0.13.1
torchvision                   0.13.1+cu113
tornado                       5.1.1
tqdm                          4.64.1
traitlets                     5.1.1
transformers                  4.22.2
tweepy                        3.10.0
typeguard                     2.7.1
typer                         0.4.2
typing-extensions             4.1.1
tzlocal                       1.5.1
ujson                         5.5.0
uritemplate                   3.0.1
urllib3                       1.24.3
vega-datasets                 0.9.0
wasabi                        0.10.1
wcwidth                       0.2.5
webargs                       8.2.0
webencodings                  0.5.1
Werkzeug                      1.0.1
wheel                         0.37.1
widgetsnbextension            3.6.1
wordcloud                     1.8.2.2
wrapt                         1.14.1
xarray                        0.20.2
xarray-einstats               0.2.2
xgboost                       0.90
xkit                          0.0.0
xlrd                          1.1.0
xlwt                          1.3.0
yarl                          1.8.1
yellowbrick                   1.5
zict                          2.2.0
zipp                          3.8.1

你可能感兴趣的:(深度学习,pytorch,python)