基于 Quanto 和 Diffusers 的内存高效 transformer 扩散模型

文章预览

过去的几个月，我们目睹了使用基于 transformer 模型作为扩散模型的主干网络来进行高分辨率文生图 (text-to-image，T2I) 的趋势。和一开始的许多扩散模型普遍使用 UNet 架构不同，这些模型使用 transformer 架构作为扩散过程的主模型。由于 transformer 的性质，这些主干网络表现出了良好的可扩展性，模型参数量可从 0.6B 扩展至 8B。 随着模型越变越大，内存需求也随之增加。对扩散模型而言，这个问题愈加严重，因为扩散流水线通常由多个模型串成: 文本编码器、扩散主干模型和图像解码器。此外，最新的扩散流水线通常使用多个文本编码器 - 如: Stable Diffusion 3 有 3 个文本编码器。使用 FP16 精度对 SD3 进行推理需要 18.765GB 的 GPU 显存。 这么高的内存要求使得很难将这些模型运行在消费级 GPU 上，因而减缓了技术采纳速度并使针对这些模型的实验变得更加困难 ………………………………