Diffusers documentation

批量推理

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

批量推理

批量推理一次处理多个提示以提升吞吐量。一个批次能让 GPU 保持忙碌,而单个提示则会导致其利用率不足。

缺点在于延迟更高,因为你需要等待整个批次处理完毕。此外,较大的批次还需要更多的 GPU 内存。

批量处理多个提示

将提示词列表传递给 pipeline,为每个提示词生成一张图像。

import torch
import matplotlib.pyplot as plt
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

prompts = [
    "Cinematic shot of a cozy coffee shop interior, warm pastel light streaming through a window where a cat rests. Shallow depth of field, glowing cups in soft focus, dreamy lofi-inspired mood, nostalgic tones, framed like a quiet film scene.",
    "Polaroid-style photograph of a cozy coffee shop interior, bathed in warm pastel light. A cat sits on the windowsill near steaming mugs. Soft, slightly faded tones and dreamy blur evoke nostalgia, a lofi mood, and the intimate, imperfect charm of instant film.",
    "Soft watercolor illustration of a cozy coffee shop interior, pastel washes of color filling the space. A cat rests peacefully on the windowsill as warm light glows through. Gentle brushstrokes create a dreamy, lofi-inspired atmosphere with whimsical textures and nostalgic calm.",
    "Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the nostalgic, lofi-inspired game aesthetic."
]

images = pipeline(
    prompt=prompts,
).images

fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes = axes.flatten()

for i, image in enumerate(images):
    axes[i].imshow(image)
    axes[i].set_title(f"Image {i+1}")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

为每个提示词生成多张图像

要为一个提示词生成多个变体,请使用 num_images_per_prompt 参数。

import torch
import matplotlib.pyplot as plt
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

prompt="""
Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the
space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the
nostalgic, lofi-inspired game aesthetic.
"""

images = pipeline(
    prompt=prompt,
    num_images_per_prompt=4
).images

fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes = axes.flatten()

for i, image in enumerate(images):
    axes[i].imshow(image)
    axes[i].set_title(f"Image {i+1}")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

为多个提示词生成多张图像

结合这两种方法,可以为不同的提示词生成不同的变体。此示例复用了上述示例中的 pipeline 和 prompts。

images = pipeline(
    prompt=prompts,
    num_images_per_prompt=2,
).images

fig, axes = plt.subplots(2, 4, figsize=(12, 12))
axes = axes.flatten()

for i, image in enumerate(images):
    axes[i].imshow(image)
    axes[i].set_title(f"Image {i+1}")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

生成可复现的批次

要使批次可复现,请向 pipeline 传入一个 Generator 对象列表,每个图像对应一个对象,且每个对象拥有各自的种子。

请参阅 可复现性 文档,以了解确定性算法和 Generator 对象。

使用列表推导式为每张图片创建一个 Generator。对列表进行乘法运算([generator] * 3)会得到指向同一个 Generator 的三个引用,因此图片无法获得独立的随机种子。

# don't do this
generator = [torch.Generator(device="cuda").manual_seed(0)] * 3  # or "mps", "xpu", "cpu"

生成器的数量必须与批次中的图像数量相匹配。将generator传递给管道。

import torch
import matplotlib.pyplot as plt
from diffusers import DiffusionPipeline

pipeline = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0",
    dtype=torch.float16,
    device_map="cuda"
)

prompts = [
    "Cinematic shot of a cozy coffee shop interior, warm pastel light streaming through a window where a cat rests. Shallow depth of field, glowing cups in soft focus, dreamy lofi-inspired mood, nostalgic tones, framed like a quiet film scene.",
    "Polaroid-style photograph of a cozy coffee shop interior, bathed in warm pastel light. A cat sits on the windowsill near steaming mugs. Soft, slightly faded tones and dreamy blur evoke nostalgia, a lofi mood, and the intimate, imperfect charm of instant film.",
    "Soft watercolor illustration of a cozy coffee shop interior, pastel washes of color filling the space. A cat rests peacefully on the windowsill as warm light glows through. Gentle brushstrokes create a dreamy, lofi-inspired atmosphere with whimsical textures and nostalgic calm.",
    "Isometric pixel-art illustration of a cozy coffee shop interior in detailed 8-bit style. Warm pastel light fills the space as a cat rests on the windowsill. Blocky furniture and tiny mugs add charm, low-res retro graphics enhance the nostalgic, lofi-inspired game aesthetic."
]
generator = [torch.Generator(device="cuda").manual_seed(i) for i in range(len(prompts))]  # or "mps", "xpu", "cpu"

images = pipeline(
    prompt=prompts,
    generator=generator
).images

fig, axes = plt.subplots(2, 2, figsize=(12, 12))
axes = axes.flatten()

for i, image in enumerate(images):
    axes[i].imshow(image)
    axes[i].set_title(f"Image {i+1}")
    axes[i].axis('off')

plt.tight_layout()
plt.show()

每张图片都关联一个种子,因此您可以选择一张喜欢的图片,在优化其提示词的同时,使用相同的种子重新生成它。

处理并发请求

服务器还能通过同时处理来自多个用户的请求来提高吞吐量。只需加载一次 pipeline 并在各请求间共享,这样模型在 GPU 内存中只需保留一份。由于调度器不是线程安全的,请使用 from_pipe() 为每个请求创建独立的调度器。在多个线程间共享同一个调度器会引发诸如 IndexError: index 21 is out of bounds for dimension 0 with size 21 之类的错误。

scheduler = shared_pipeline.scheduler.from_config(shared_pipeline.scheduler.config)
pipeline = StableDiffusion3Pipeline.from_pipe(shared_pipeline, scheduler=scheduler)

对于使用 FastAPI 构建的完整服务器,请参阅 server 示例和 server-async 示例。server-async 示例还会为每个请求克隆随机数生成器的状态,并使分词器(tokenizer)可以安全地在多个线程间共享。

后续步骤

  • 当单个批次无法放入一块 GPU,或者提示词数量超过单块 GPU 的处理能力时,请使用数据并行将提示词分布到多块 GPU 上。
Update on GitHub