在低VRAM GPU上运行 PixArt-Σ/Flux.1 图像生成：Python 短教程

7 个月前

在低显存 GPU 上运行 PixArt-Σ/Flux.1 图像生成：Python 简短教程

Diffusers 和 Quanto 为显存不足的 GPU 带来了希望

由 PixArt-Σ 在少于 8GB 显存的情况下本地生成

图像生成工具正如火如荼，功能也越来越强大。像 PixArt Sigma 和 Flux.1 这样的模型引领了潮流，得益于它们开放的权重模型和宽松的许可证。这种设置允许进行创造性的尝试，包括在不共享数据的情况下训练 LoRAs。

然而，如果你使用的是较旧或显存较少的 GPU，使用这些模型可能会面临挑战。通常，质量、速度和显存使用之间存在权衡。在这篇博客中，我们将重点优化速度和显存使用，同时尽可能保持质量。这种方法对于 PixArt 特别有效，因为它的体积较小，但在 Flux.1 上的结果可能会有所不同。我将在本文末尾分享一些 Flux.1 的替代解决方案。

PixArt Sigma 和 Flux.1 都是基于变换器的，这意味着它们受益于大型语言模型（LLMs）使用的相同量化技术。量化涉及压缩模型的组件以减少内存使用。这使你能够同时将所有模型组件保留在 GPU 显存中，从而比在 GPU 和 CPU 之间移动权重的方法更快。

让我们开始设置吧！

设置本地环境

首先，确保你安装了 Nvidia 驱动程序和 Anaconda。

接下来，创建一个 Python 环境并安装所有主要依赖项：

conda install pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia

然后安装 Diffusers 和 Quanto 库：

pip install pillow==10.3.0 loguru~=0.7.2 optimum-quanto==0.2.4 diffusers==0.30.0 transformers==4.44.2 accelerate==0.33.0 sentencepiece==0.2.0

量化代码

以下是一个简单的脚本，帮助你开始使用 PixArt-Sigma：

from optimum.quanto import qint8, qint4, quantize, freeze
from diffusers import PixArtSigmaPipeline
import torch

pipeline = PixArtSigmaPipeline.from_pretrained(
    "PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
)

quantize(pipeline.transformer, weights=qint8)
freeze(pipeline.transformer)

quantize(pipeline.text_encoder, weights=qint4, exclude="proj_out")
freeze(pipeline.text_encoder)

pipe = pipeline.to("cuda")

for i in range(2):
    generator = torch.Generator(device="cpu").manual_seed(i)

    prompt = "赛博朋克城市风景，小黑乌鸦，霓虹灯，黑暗小巷，摩天大楼，未来主义，鲜艳的色彩，高对比度，高度细致"

    image = pipe(prompt, height=512, width=768, guidance_scale=3.5, generator=generator).images[0]

    image.save(f"Sigma_{i}.png")

理解脚本：实现的主要步骤

  1. 导入必要的库：我们导入用于量化、模型加载和 GPU 处理的库。
  2. 加载模型：我们将 PixArt Sigma 模型以半精度（float16）加载到 CPU。
  3. 量化模型：我们对模型的变换器和文本编码器组件进行量化。这里我们应用不同级别的量化：文本编码器部分使用 qint4 进行量化，因为它相对较大。视觉部分如果使用 qint8 进行量化，整个管道将使用 7.5 GB 显存，如果不进行量化则大约使用 8.5 GB 显存。
  4. 移动到 GPU：我们将管道移动到 GPU .to("cuda") 以加快处理速度。
  5. 生成图像：我们使用 pipe 根据给定的提示生成图像并保存输出。

运行脚本

保存脚本并在你的环境中运行。你应该会看到根据提示“赛博朋克城市风景，小黑乌鸦，霓虹灯，黑暗小巷，摩天大楼，未来主义，鲜艳的色彩，高对比度，高度细致”生成的图像，保存为 sigma_1.png。在 RTX 3080 GPU 上生成时间为 6 秒。

由 PixArt-Σ 本地生成

尽管 Flux.1 Schnell 有额外的组件，但你仍然可以获得类似的结果，但这需要更激进的量化，这可能会降低质量（除非你有更多的显存，比如 16 或 25 GB）。

import torch

from optimum.quanto import qint2, qint4, quantize, freeze

from diffusers.pipelines.flux.pipeline_flux import FluxPipeline

pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)

quantize(pipe.text_encoder, weights=qint4, exclude="proj_out")
freeze(pipe.text_encoder)

quantize(pipe.text_encoder_2, weights=qint2, exclude="proj_out")
freeze(pipe.text_encoder_2)

quantize(pipe.transformer, weights=qint4, exclude="proj_out")
freeze(pipe.transformer)

pipe = pipe.to("cuda")

for i in range(10):
    generator = torch.Generator(device="cpu").manual_seed(i)
    prompt = "赛博朋克城市风景，小黑乌鸦，霓虹灯，黑暗小巷，摩天大楼，未来主义，鲜艳的色彩，高对比度，高度细致"

    image = pipe(prompt, height=512, width=768, guidance_scale=3.5, generator=generator, num_inference_steps=4).images[0]

    image.save(f"Schnell_{i}.png")

由 Flux.1 Schnell 本地生成：由于过度量化，质量较低且对提示的遵循性差

我们可以看到，将文本编码器量化为 qint2 和视觉变换器量化为 qint8 可能过于激进，这对 Flux.1 Schnell 的质量产生了显著影响。

以下是运行 Flux.1 Schnell 的一些替代方案：

如果 PixArt-Sigma 不能满足你的需求，并且你没有足够的显存以足够的质量运行 Flux.1，你有两个主要选择：

  • ComfyUI 或 Forge：这些是爱好者使用的 GUI 工具，通常牺牲速度以换取质量。
  • Replicate API：每生成一张图像的费用为 0.003。

部署

我在一台旧机器上部署 PixArt Sigma 时玩得很开心。以下是我进行的简要总结：

首先是组件列表：

  1. HTMX 和 Tailwind：这些就像项目的面子。HTMX 帮助使网站互动而不需要很多额外代码，Tailwind 则让它看起来不错。
  2. FastAPI：它接收来自网站的请求并决定如何处理。
  3. Celery Worker：可以把它看作是辛勤的工作者。它从 FastAPI 接收订单并实际创建图像。
  4. Redis 缓存/发布-订阅：这就像通信中心。它帮助项目的不同部分相互交流并记住重要信息。
  5. GCS（Google Cloud Storage）：这是我们存放完成图像的地方。

那么，它们是如何协同工作的呢？以下是简单的流程：

  • 当你访问网站并发出请求时，HTMX 和 Tailwind 确保它看起来不错。
  • FastAPI 获取请求并通过 Redis 告诉 Celery Worker 生成什么样的图像。
  • Celery Worker 开始工作，创建图像。
  • 一旦图像准备好，它会存储在 GCS 中，以便于访问。

应用程序演示

结论

通过量化模型组件，我们可以显著减少显存使用，同时保持良好的图像质量并提高生成速度。这种方法对于像 PixArt Sigma 这样的模型特别有效。对于 Flux.1，尽管结果可能不尽如人意，但量化的原则仍然适用。

参考文献：

FluxAI 中文

© 2025. All Rights Reserved

FLUX.1 by Black Forest Labs