基于Tensorflow、Keras实现Stable Diffusion,开箱即用实现多GPU推理

2022-09-21 16:42:59 浏览数 (2)

机器之心报道

机器之心编辑部

研究者对 Stable Diffusion 的探索,永无止境。

最近一段时间,文本转图像模型 Stable Diffusion 可谓是爆红 AI 圈,其是由慕尼黑大学和 Runway 的研究者基于 CVPR 2022 的论文《High-Resolution Image Synthesis with Latent Diffusion Models》实现的,它可以在消费级 GPU 上运行。

自推出以来,研究者纷纷对 Stable Diffusion 进行了更深入的研究。有研究者将其和 Web UI「拼在」一起,两者组合成绘画工具,让没有系统学习过 UI 知识的小伙伴,也可以上手操作。

为了帮助更多的人快速上手,还有研究者整理了一份在苹果 M1 芯片上运行 Stable Diffusion 的指南。

不过原有模型是基于 Torch 实现的,现在,来自 Meta 的研究者 Divam Gupta 表示:基于 Tensorflow/Keras 实现的 Stable Diffusion 已经来了。

项目地址:https://github.com/divamgupta/stable-diffusion-tensorflow

总结来说,该项目的特点有 3 个:转换预训练模型;代码易于理解;代码量少。

Divam Gupta 表示,在 8GB M1 MacBook Air 上运行 Stable Diffusion 时,与 Torch 实现相比,运行速度大约快 4 倍。

Keras 的创造者 François Chollet 表示:它在 M1 MacBooPros GPU 上实现开箱即用,它还可以开箱即用地进行多 GPU 推理。

Chollet 还表示,除了GPU外,它还可以开箱即用地进行 TPU 推理:只需获取一个 TPU VM 并将 TPU strategy scope 添加到代码中。在进行大批量处理时,这可以显着加快推理速度、降低成本。

最后,Chollet 表示「非常感谢 Divam Gupta 创建这个端口。这是一项高质量的工作,将使所有从事创造性人工智能的人受益。我总是惊讶于开源社区的发展速度。」

项目介绍

怎样运行这个项目呢?

先进入这个地址:https://colab.research.google.com/drive/1zVTa4mLeM_w44WaFwl7utTaa6JcaH1zK

运行命令行:

代码语言:javascript复制
python text2image.py --prompt="An astronaut riding a horse"

使用 python 接口:

代码语言:javascript复制
pip install git https://github.com/fchollet/stable-diffusion-tensorflow
代码语言:javascript复制
from stable_diffusion_tf.stable_diffusion import Text2Imagefrom PIL import Image
generator = Text2Image(     img_height=512,    img_width=512,    jit_compile=False,)img = generator.generate(    "An astronaut riding a horse",    num_steps=50,    unconditional_guidance_scale=7.5,    temperature=1,  batch_size=1,)Image.fromarray(img[0]).save("output.png")

简单两步就能完成,输出结果如下:一位洛可可风格的狼人在燃烧的咖啡店里喝咖啡。

摩天大楼上粉白的蜘蛛侠:

天堂中的景象:

有兴趣的小伙伴,可以上手一试。

参考链接:

https://twitter.com/divamgupta/status/1571234504320208897?s=20&t=kpyjZ8LLSmgJFnUKqyepAw

https://twitter.com/fchollet/status/1571874757582389250

© THE END 

转载请联系本公众号获得授权

投稿或寻求报道:content@jiqizhixin.com

0 人点赞