LlamaFactory
算力平台:

JAX/Flax ​

[[open-in-colab]]

🤗 Diffusers 支持 Flax，可以在 Google TPU 上进行超快的推理，例如在 Colab、Kaggle 或 Google Cloud Platform 上可用的 TPU。本指南将向你展示如何使用 JAX/Flax 运行 Stable Diffusion 的推理。

在开始之前，请确保你已经安装了必要的库：

py
# uncomment to install the necessary libraries in Colab
#!pip install -q jax==0.3.25 jaxlib==0.3.25 flax transformers ftfy
#!pip install -q diffusers

你还应该确保使用的是 TPU 后端。虽然 JAX 并不只在 TPU 上运行，但你可以在 TPU 上获得最佳性能，因为每个服务器都有 8 个 TPU 加速器并行工作。

如果你在 Colab 中运行本指南，请在上方菜单中选择 Runtime，然后选择 Change runtime type，在 Hardware accelerator 设置中选择 TPU。导入 JAX 并快速检查你是否在使用 TPU：

python
import jax
import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

num_devices = jax.device_count()
device_type = jax.devices()[0].device_kind

print(f"Found {num_devices} JAX devices of type {device_type}.")
assert (
    "TPU" in device_type,
    "Available device is not a TPU, please select TPU from Runtime > Change runtime type > Hardware accelerator"
)
# Found 8 JAX devices of type Cloud TPU.

好的，现在你可以导入其余所需的依赖项：

python
import jax.numpy as jnp
from jax import pmap
from flax.jax_utils import replicate
from flax.training.common_utils import shard

from diffusers import FlaxStableDiffusionPipeline

加载模型 ​

Flax 是一个函数式框架，因此模型是无状态的，参数存储在模型之外。加载预训练的 Flax 管道会返回 模型和模型权重（或参数）。在本指南中，你将使用 bfloat16，这是一种更高效的半精度浮点类型，受 TPUs 支持（你也可以使用 float32 以获得全精度）。

python
dtype = jnp.bfloat16
pipeline, params = FlaxStableDiffusionPipeline.from_pretrained(
    "CompVis/stable-diffusion-v1-4",
    variant="bf16",
    dtype=dtype,
)

推理 ​

TPU 通常有 8 个设备并行工作，因此我们可以为每个设备使用相同的提示。这意味着你可以在 8 个设备上同时进行推理，每个设备生成一张图像。因此，你可以在与单个芯片生成一张图像相同的时间内获得 8 张图像！

在复制提示后，通过调用管道上的 prepare_inputs 函数获取分词后的文本 ID。分词后的文本长度设置为 77 个 tokens，这是底层 CLIP 文本模型的配置要求。

python
prompt = "A cinematic film still of Morgan Freeman starring as Jimi Hendrix, portrait, 40mm lens, shallow depth of field, close up, split lighting, cinematic"
prompt = [prompt] * jax.device_count()
prompt_ids = pipeline.prepare_inputs(prompt)
prompt_ids.shape
# (8, 77)

模型参数和输入必须在8个并行设备上进行复制。参数字典使用 flax.jax_utils.replicate 进行复制，该函数会遍历字典并改变权重的形状，使它们重复8次。数组使用 shard 进行复制。

python
# parameters
p_params = replicate(params)

# arrays
prompt_ids = shard(prompt_ids)
prompt_ids.shape
# (8, 1, 77)

这个形状表示每个设备接收一个形状为 (1, 77) 的 jnp 数组，其中 1 是每个设备的批量大小。在具有足够内存的 TPUs 上，如果你希望一次生成多个图像（每个芯片），你可以使用大于 1 的批量大小。

接下来，创建一个随机数生成器以传递给生成函数。这是 Flax 中的标准做法，Flax 对随机数的处理非常严格和有主见。所有处理随机数的函数都期望接收一个生成器，以确保可重复性，即使在多个分布式设备上进行训练时也是如此。

下面的辅助函数使用一个种子来初始化随机数生成器。只要你使用相同的种子，你将获得完全相同的结果。在后续的指南中探索结果时，可以自由使用不同的种子。

python
def create_key(seed=0):
    return jax.random.PRNGKey(seed)

辅助函数，或 rng，被拆分 8 次，以便每个设备接收一个不同的生成器并生成不同的图像。

python
rng = create_key(0)
rng = jax.random.split(rng, jax.device_count())

为了利用 JAX 在 TPU 上的优化速度，可以在管道中传递 jit=True，将 JAX 代码编译成高效的表示形式，并确保模型在 8 个设备上并行运行。

第一次推理运行需要更多时间，因为它需要编译代码，但后续调用（即使输入不同）会快得多。例如，在 TPU v2-8 上编译需要超过一分钟，但之后的推理运行只需要大约 7秒！

py
%%time
images = pipeline(prompt_ids, p_params, rng, jit=True)[0]

# CPU times: user 56.2 s, sys: 42.5 s, total: 1min 38s
# Wall time: 1min 29s

返回的数组形状为 (8, 1, 512, 512, 3)，需要重塑以移除第二个维度，从而得到 8 张 512 × 512 × 3 的图像。然后你可以使用 [~utils.numpy_to_pil] 函数将数组转换为图像。

python
from diffusers.utils import make_image_grid

images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)
make_image_grid(images, rows=2, cols=4)

img

使用不同的提示 ​

你不必在所有设备上使用相同的提示。例如，生成 8 个不同的提示：

python
prompts = [
    "Labrador in the style of Hokusai",
    "Painting of a squirrel skating in New York",
    "HAL-9000 in the style of Van Gogh",
    "Times Square under water, with fish and a dolphin swimming around",
    "Ancient Roman fresco showing a man working on his laptop",
    "Close-up photograph of young black woman against urban background, high quality, bokeh",
    "Armchair in the shape of an avocado",
    "Clown astronaut in space, with Earth in the background",
]

prompt_ids = pipeline.prepare_inputs(prompts)
prompt_ids = shard(prompt_ids)

images = pipeline(prompt_ids, p_params, rng, jit=True).images
images = images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:])
images = pipeline.numpy_to_pil(images)

make_image_grid(images, 2, 4)

img

并行化是如何工作的？ ​

🤗 Diffusers 中的 Flax 管道会自动编译模型并在所有可用设备上并行运行。让我们更详细地看看这个过程是如何工作的。

JAX 并行化可以通过多种方式实现。最简单的方法是使用 jax.pmap 函数来实现单程序多数据（SPMD）并行化。这意味着运行多个相同代码的副本，每个副本处理不同的数据输入。更复杂的实现方式也是可能的，如果你感兴趣，可以查阅 JAX 文档 以深入了解这个主题！

jax.pmap 做了两件事：

  1. 编译（或“jit”）代码，类似于 jax.jit()。这不会在调用 pmap 时发生，而是在第一次调用 pmapped 函数时发生。
  2. 确保编译后的代码在所有可用设备上并行运行。

为了演示，我们可以在管道的 _generate 方法上调用 pmap（这是一个生成图像的私有方法，可能会在未来的 🤗 Diffusers 版本中重命名或移除）：

python
p_generate = pmap(pipeline._generate)

调用 pmap 后，准备好的函数 p_generate 将会：

  1. 在每个设备上复制底层函数 pipeline._generate。
  2. 将输入参数的不同部分发送到每个设备（这就是为什么需要调用 shard 函数）。在这种情况下，prompt_ids 的形状为 (8, 1, 77, 768)，因此数组被分成 8 份，每个 _generate 的副本接收一个形状为 (1, 77, 768) 的输入。

这里最重要的是要注意批量大小（在这个例子中为 1）和适合你代码的输入维度。你不需要做其他任何更改来使代码并行运行。

第一次调用管道会花费更多时间，但之后的调用会快得多。block_until_ready 函数用于正确测量推理时间，因为 JAX 使用异步调度，并在可以时立即将控制权返回给 Python 循环。你不需要在代码中使用它；当你需要使用尚未实现的计算结果时，阻塞会自动发生。

py
%%time
images = p_generate(prompt_ids, p_params, rng)
images = images.block_until_ready()

# CPU times: user 1min 15s, sys: 18.2 s, total: 1min 34s
# Wall time: 1min 15s

检查你的图片尺寸，看看它们是否正确：

python
images.shape
# (8, 1, 512, 512, 3)

资源 ​

要了解更多关于 JAX 如何与 Stable Diffusion 配合使用，你可能对以下内容感兴趣：

Copyright © 2024-2025 llamafactory.cn