如何在StableDiffusionXLImg2ImgPipeline中实现批量图生图

针对批量生成与文件夹内图片相似图像的需求，HuggingFace生态里有几种高效的实现方式，无需手动单张循环或重复初始化Pipeline：

1. 利用Diffusers Pipeline的原生批量输入支持 #

Diffusers的StableDiffusionXLImg2ImgPipeline本身支持接收批量初始化图像，GPU会自动并行处理，比单张循环效率更高。

代码示例 #

import torch
import os
from diffusers.utils import load_image
from diffusers import StableDiffusionXLImg2ImgPipeline

# 初始化并加载Pipeline（只需执行一次）
pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0", 
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True
)
pipe = pipe.to("cuda")

# 批量加载目标文件夹内的图像
image_dir = "MyImages\\ImageList\\"
# 过滤有效图像格式
image_paths = [
    os.path.join(image_dir, filename) 
    for filename in os.listdir(image_dir) 
    if filename.lower().endswith((".jpg", ".png", ".jpeg"))
]
# 加载并统一预处理图像
init_images = [load_image(path).convert("RGB") for path in image_paths]

# 批量生成相似图像（支持单prompt或对应每个图像的prompt列表）
prompt = "Generate an image similar to the input"
# 若需为不同图像指定不同prompt，可传入列表：prompt = ["prompt1", "prompt2", ...]
generated_images = pipe(prompt, image=init_images).images

# 批量保存结果
for idx, img in enumerate(generated_images):
    img.save(f"generated_similar_{idx}.png")

2. 分批次处理适配显存限制 #

如果图像数量多、显存不足，可以按固定批次大小拆分处理，避免显存溢出：

代码示例 #

batch_size = 4  # 根据GPU显存调整，比如12GB显存可设为2-4

for batch_start in range(0, len(init_images), batch_size):
    # 截取当前批次的图像
    current_batch = init_images[batch_start:batch_start+batch_size]
    # 生成当前批次的图像
    batch_results = pipe(prompt, image=current_batch).images
    # 保存当前批次结果
    for offset, img in enumerate(batch_results):
        img.save(f"generated_similar_{batch_start + offset}.png")

3. 用HuggingFace Datasets管理大规模批量任务 #

如果处理的是大规模图像数据集，推荐用datasets库来简化数据加载、预处理和批量工作流，支持多进程预处理、缓存等功能：

代码示例 #

from datasets import load_dataset

# 加载本地图像文件夹为数据集
dataset = load_dataset("imagefolder", data_dir=image_dir)

# 定义生成函数
def generate_similar(example):
    # 预处理输入图像
    init_img = example["image"].convert("RGB")
    # 生成相似图像
    generated_img = pipe(prompt, image=init_img).images[0]
    example["generated_image"] = generated_img
    return example

# 批量处理数据集（batched=True可配合Pipeline批量输入进一步优化）
processed_dataset = dataset.map(generate_similar, batched=False)

# 保存所有生成的图像
for idx, item in enumerate(processed_dataset["train"]):
    item["generated_image"].save(f"generated_similar_{idx}.png")

关键说明 #

• 所有方案均无需重复初始化Pipeline，只需加载一次即可复用，大幅节省时间。
• 批量处理的核心是利用GPU的并行计算能力，比单张循环效率提升明显。
• 若需进一步优化性能，可结合accelerate库实现分布式训练/推理，适合多GPU场景。

内容的提问来源于stack exchange，提问作者Adarsh Wase