批量推理
批量推理一次处理多个提示以提升吞吐量。一个批次能让 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 上。