StableDiffusion图像生成能力一探！附Int8量化教程与ONNX导出推理

2022-09-19 10::27

只吃瓜怎么能行，当然要上手一试啦！本文送上了一份StableFusion的int8量化教程，以及ONNX导出，甚至是TensorRT加速推理。通过CPU OpenVINO加速也就3分钟推理。

最近stablediffusion大火，但很多人都只是吃瓜，最多也就是在huggingface网站上试一下，这其实并不够，作为一个富有商业嗅觉的AI从业者，我嗅探到的更多的是他的商业能力，不得不说，现在生成类的AI模型，已经越来越来接近甚至超越人类水平。

今天就来一探究竟，看看StableFusion到底能做啥？并附上这个超级大模型的int8量化教程，以及ONNX导出，甚至是TensorRT加速推理。

请注意，一下的所有测试案例，都是我用CPU跑出来的，你要问我多慢？通过CPU OpenVINO加速也就3分钟推理。为什么不用GPU？因为显存不足！

好在我们用CPU也能比较完善的推理，结果并不差。未来我们还会继续使用TensorRT推理，加速整个推理流程。更进一步的将StableDiffusion变成生产力工具！

代码

这次我们直接先上代码，这里面包含了CPU下推理StableDiffusion，以及OpenVINO加速的代码，同时，也包含了量化脚本，感觉兴趣的可以再github提issue。

git clone https://github.com/luohao123/gaintmodels

备注：代码来自大佬，不是本人写的。不过有问题欢迎提问交流。

测试StableDiffusion

来看看生成的效果，由于模型只能编码英文，我们就以英文作为promopt。

A green car with appearance of Tesla Model 3 and Porsche 911

这长得还真的就像是，保时捷和model3的合体！中间那个logo你注意看，感觉并不是特斯拉的logo！有点像保时捷的logo！~

A robot Elon Musk in cyberpunk, driving on a Tesla Model X

ModelX倒是有点像，但是马斯克人呢？？

StableDiffusion这类模型，似乎在生成风景画上表现不错

A beautiful mansion beside a lake in the woods, with a clean road front of it.

一栋湖边的别墅

A beautiful castle beside a waterfall in the woods, detailed, 4k

瀑布旁边的城堡

真的是非常的优美！而且细节清晰可见！！

接下来我们测试一下，一些NSFW的内容，请注意，高能预警，我们的模型把huggingface的SaftyChecker去掉了，为什么？因为我们并不需要！并且这个模块是会占用很多内存的！

A beautiful sexy girl with red hair, 4k, detailed, without any cloth.

由于你懂得原因，我只截取了一部分！！但是，这个生成的效果真的令人震惊！过于真实！！

再来测试一些创造性的东西：

A advanced spaceship with Elon Musk driving on it, detailed, in reality style

这个非常设计的，地区非常的逼真！看这个反光，这个倒影，以及上面的mount，甚至有种自动驾驶飞船的感觉。

A red tv in front of sofa, a child is looking at it with a dog.

在比如这个图片，这个真实度简直超乎了我的预料，电视机里面甚至还可以看到倒影！！

代码讲解

Experiements on testing GaintModels such as GPT3, StableFusion. We offer TensorRT && Int8 quantization on those gaint models. Make you can inference on a 6GB below GPU mem card!

Install

Some requirements to install:

pip install diffusers
pip install transformers
pip install alfred-py

Models

StableFusion:

First, we need download stablefusion weights from hugging face.

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

You should downloading weights using git lfs large file system, the model about 3GB.

To make unet_2d_condition in stablefusion able to export to onnx, make some modification on diffusers, following: link

file: diffuers/models/unet_2d_conditions.py

# L137
timesteps = timesteps.broadcast_to(sample.shape[0])
#timesteps = timesteps.broadcast_to(sample.shape[0])
timesteps = timesteps * torch.ones(sample.shape[0])

output = {"sample": sample}
#output = {"sample": sample}

return output
return sample

After that, move stable-diffusion-v1-4 to weights folder. Run: