当前位置:AIGC资讯 > AIGC > 正文

Stable Diffusion XL on diffusers

Stable Diffusion XL on diffusers

翻译自:https://huggingface.co/docs/diffusers/using-diffusers/sdxl v0.24.0 非逐字翻译

Stable Diffusion XL (SDXL) 是一个强大的图像生成模型,其在上一代 Stable Diffusion 的基础上主要做了如下优化:

参数量增加:SDXL 中 Unet 的参数量比前一代大了 3 倍,并且 SDXL 还引入了第二个 text-encoder(OpenCLIP ViT-bigG/14),整体参数量大幅增加。 引入了 size-conditioning 和 crop conditioning,在训练阶段有效利用起低分辨率图像,并在推理对生成的图片是否需要裁剪有更好的控制。 使用了两阶段的生成过程,除了第一阶段的 base 模型生成之外,还加入了一个 refiner 模型,使得生成图像的细节质量更高(其中 base 模型也可以单独使用,直接生成)

本文将介绍如何使用 diffusers 进行 text-to-image、image-to-image 和 inpainting。

加载模型参数

模型参数分别保存在不同的子目录中,可以使用 from_pretrained 方法来加载:

from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
import torch

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

refiner = StableDiffusionXLImg2ImgPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0", torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
).to("cuda")

也可以使用 from_single_file 方法来从单个文件 ( .ckpt 或 .safetensors) 中加载:

from diffusers import StableDiffusionXLPipeline, StableDiffusionXLImg2ImgPipeline
import torch

pipeline = StableDiffusionXLPipeline.from_single_file(
    "https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0/blob/main/sd_xl_base_1.0.safetensors", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

refiner = StableDiffusionXLImg2ImgPipeline.from_single_file(
    "https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0/blob/main/sd_xl_refiner_1.0.safetensors", torch_dtype=torch.float16, use_safetensors=True, variant="fp16"
).to("cuda")

text-to-image

在进行 text-to-image 生成时,需要传入文本 prompt。SDXL 默认生成分辨率为 1024 * 1024,也可以设置为 768 或 512,但不要再低于 512 了:

from diffusers import AutoPipelineForText2Image
import torch

pipeline_text2image = AutoPipelineForText2Image.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipeline_text2image(prompt=prompt).images[0]
image

image-to-image

在进行 image-to-image 生成时,SDXL 在 768 - 1024 这个分辨率区间工作的最好。此时需要输入一张原始图像,并给一段文本 prompt:

from diffusers import AutoPipelineForImage2Image
from diffusers.utils import load_image, make_image_grid

# 使用 from_pipe,避免在加载 checkpoint 时消耗额外的内存
pipeline = AutoPipelineForImage2Image.from_pipe(pipeline_text2image).to("cuda")

url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
init_image = load_image(url)
prompt = "a dog catching a frisbee in the jungle"
image = pipeline(prompt, image=init_image, strength=0.8, guidance_scale=10.5).images[0]
make_image_grid([init_image, image], rows=1, cols=2)

inpainting

在记性 inpainting 时,需要传入一张原始图片和原始图片中你想要修改部分的 mask 图,并给一个文本 prompt 来描述 mask 区域需要生成什么内容:

from diffusers import AutoPipelineForInpainting
from diffusers.utils import load_image, make_image_grid


pipeline = AutoPipelineForInpainting.from_pipe(pipeline_text2image).to("cuda")

img_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-text2img.png"
mask_url = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/sdxl-inpaint-mask.png"

init_image = load_image(img_url)
mask_image = load_image(mask_url)

prompt = "A deep sea diver floating"
image = pipeline(prompt=prompt, image=init_image, mask_image=mask_image, strength=0.85, guidance_scale=12.5).images[0]
make_image_grid([init_image, mask_image, image], rows=1, cols=3)

refine image quality

SDXL 相比于之前的 SD 模型,一个很大的差别在于它包含了一个 refiner 模型。refiner 模型 可以接收 base 模型经过几步去噪之后的低噪声图像,并为图像生成更多高质量的细节。有两种使用 refiner 模型的方式:

同时使用 base 模型和 refiner 模型,来生成高质量图片 使用 base 模型生成一张图片,然后使用 refiner 模型为图片添加更多的细节(这是 SDXL 训练时的方式)

接下来分别介绍这两种方式的使用。

base + refiner model

当使用第一种方式,即同时使用 base 模型和 refiner 模型来生成图片时,称为 ensemble of expert denoisers。这种方式相比于第二种将 base 模型的输出给 refiner 模型中的方式来说,整体需要的去噪步数更少,因此会快很多。但这种方式我们看到的 base 模型的输出是带有一些噪声的。

在第一种方式中,base 模型负责高噪声阶段的去噪,refiner 模型负责低噪声阶段的去噪。首先加载 base 模型和 refiner 模型:

from diffusers import DiffusionPipeline
import torch

base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
).to("cuda")

在使用 ensemble of expert denoisers 这种方式时,我们需要定义不同的模型在他们各自阶段的去噪步数。对于 base 模型,需要 denoising_end 参数,对于 refiner 模型,需要 denoising_start 参数。

denoising_startdenoising_end 参数都时 0-1 之间的一个小数,用于表示当前 schduler 下步数的比例。如果同时还传入了 strength 参数,它将被忽略,因为去噪步骤的数量是由模型训练的离散时间步长和声明的比例截止值决定的。

这里,我们设置 denoising_end 为 0.8,从而 base 模型会负责前 80% 的高噪声阶段的降噪,并设置 denoising_start 为 0.8,从而 refiner 模型会负责后 20% 的低噪声阶段的降噪。注意 base 模型的输出是在隐层 latent 空间的,而非可见的图片。

prompt = "A majestic lion jumping from a big stone at night"

image = base(
    prompt=prompt,
    num_inference_steps=40,
    denoising_end=0.8,
    output_type="latent",
).images
image = refiner(
    prompt=prompt,
    num_inference_steps=40,
    denoising_start=0.8,
    image=image,
).images[0]
image

在 StableDiffusionXLInpaintPipeline 中,refiner 模型也可以用于进行 inpainting:

from diffusers import StableDiffusionXLInpaintPipeline
from diffusers.utils import load_image, make_image_grid
import torch

base = StableDiffusionXLInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

refiner = StableDiffusionXLInpaintPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
).to("cuda")

img_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo.png"
mask_url = "https://raw.githubusercontent.com/CompVis/latent-diffusion/main/data/inpainting_examples/overture-creations-5sI6fQgYIuo_mask.png"

init_image = load_image(img_url)
mask_image = load_image(mask_url)

prompt = "A majestic tiger sitting on a bench"
num_inference_steps = 75
high_noise_frac = 0.7

image = base(
    prompt=prompt,
    image=init_image,
    mask_image=mask_image,
    num_inference_steps=num_inference_steps,
    denoising_end=high_noise_frac,
    output_type="latent",
).images
image = refiner(
    prompt=prompt,
    image=image,
    mask_image=mask_image,
    num_inference_steps=num_inference_steps,
    denoising_start=high_noise_frac,
).images[0]
make_image_grid([init_image, mask_image, image.resize((512, 512))], rows=1, cols=3)

这种 ensemble of expert denoisers 的方式对于所有 scheduler 都可用。

base to refiner model

第二种方式通过 base 模型先生成一张完全去噪的图片,然后使用 refiner 模型以 image-to-image 的形式,为图片添加更多的高质量细节,这使得 SDXL 的生成质量有了极大的提高。首先加载 base 和 refiner 模型:

from diffusers import DiffusionPipeline
import torch

base = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

refiner = DiffusionPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-refiner-1.0",
    text_encoder_2=base.text_encoder_2,
    vae=base.vae,
    torch_dtype=torch.float16,
    use_safetensors=True,
    variant="fp16",
).to("cuda")

先使用 base 模型生成一张图片,注意将输出形式设置为 latent:

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"

image = base(prompt=prompt, output_type="latent").images[0]

将生成的图片输入到 refiner 模型中:

image = refiner(prompt=prompt, image=image[None, :]).images[0]

要进行 inpainting,在 StableDiffusionXLInpaintPipeline 中加载 base 和 refiner 模型,去掉 denoising_enddenoising_start 参数,并为 refiner 模型设置一个较小的步数。

micro-conditioning

SDXL 训练时使用了许多额外的条件方式,即 micro-conditioning,包括 original_image_size、target_image_size 和 cropping parameters。在推理阶段,合理地使用 micro-conditioning 可以生成高质量的、居中的图片。

由于 classfier-free guidance 的存在,可以在 SDXL 相关的 pipeline 中使用 micro-conditioning 和 negative micro-conditioning 参数。

size conditioning

size conditioning 有两种:

original size conditioning。训练集中有许多图片的分辨率是较低的,但又不能直接不用这些低分辨率图像(占比达 40%,丢了太浪费了),因此通常会对这些图像进行 resize,从而得到高分辨率的图像。在这个过程中,不可避免得会引入插值这种人工合成的模糊痕迹,被 SDXL 学到,而在真正的高分辨率图像中,是不该有这些痕迹的。因此训练时会告诉模型,这张图片实际是多少分辨率的,作为条件。

在推理阶段,我们可以指定 original_size 来表示图像的原始尺寸。使用默认的 1024,能生成出与原始数据集中高分辨率图像类似的高质量图像。而如果将这个值设得很低,如 256,模型还是会生成分辨率为 1024 的图像,但就会带有低分辨率图像的特征(如模糊、模式简单等)。

target size conditioning。SDXL 训练时支持多种不同的长宽比。

推理时,如果使用默认的值 1024,生成的图像会看起来像方形图像(长宽比1:1)。这里建议将 target_size 和 original_size 设置为相同的值,但你也可以调一调这些参数实验一下看看。

在 diffusers 中,我们还可以指定有关图像大小的 negative 条件,从而引导生成远离某些图像分辨率:

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(
    prompt=prompt,
    negative_original_size=(512, 512),
    negative_target_size=(1024, 1024),
).images[0]
crop conditioning

SDXL 之前的 SD 模型的生成结果有时会看起来像是被裁剪过得。这是因为为了保证训练时每个 batch 内的尺寸一致,输入的训练数据确实有很多是裁剪过的。因此训练时,裁剪坐标也会作为条件给到模型。从而,在推理时,我们将裁剪坐标指定为 (0, 0) (也是 diffusers 默认值),就可以生成非裁剪的图片了。你也可以试着调一下裁剪坐标这个参数,看模型的生成结果会是什么样子,应该可以得到非居中的构图。

from diffusers import StableDiffusionXLPipeline
import torch

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipeline(prompt=prompt, crops_coords_top_left=(256, 0)).images[0]
image

同样可以指定 negative 裁剪坐标以引导生成远离某些裁剪参数:

from diffusers import StableDiffusionXLPipeline
import torch

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
image = pipe(
    prompt=prompt,
    negative_original_size=(512, 512),
    negative_crops_coords_top_left=(0, 0),
    negative_target_size=(1024, 1024),
).images[0]
image

Use a different prompt for each text-encoder

SDXL 有两个 text encoder,所以给两个 text encoder 传入不同的文本 prompt 是可能的,这可以提高生成质量(参考)。将原本的 prompt 传到 prompt 中,另一个 prompt 传到 prompt_2 中。如果使用 negative prompt 也是类似的,分别传到 negative_promptnegative_prompt_2

from diffusers import StableDiffusionXLPipeline
import torch

pipeline = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", torch_dtype=torch.float16, variant="fp16", use_safetensors=True
).to("cuda")

# prompt is passed to OAI CLIP-ViT/L-14
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# prompt_2 is passed to OpenCLIP-ViT/bigG-14
prompt_2 = "Van Gogh painting"
image = pipeline(prompt=prompt, prompt_2=prompt_2).images[0]
image

SDXL 的双 text encoder 同样支持 textual inversion embeddings,需要分别加载,详情见:SDXL textual inversion 。

Optimizations

SDXL 的模型还是很大的,可能在一些设备上运行会比较吃力,以下是一些节约内存和提高推理速度的技巧。

如果显存不够,可以临时将模型 offload 到内存中

# base.to("cuda")
# refiner.to("cuda")
base.enable_model_cpu_offload()
refiner.enable_model_cpu_offload()

如果你使用的 torch 版本 > 2.0,那么使用 torch.cmpile 可以提速约 20%

base.unet = torch.compile(base.unet, mode="reduce-overhead", fullgraph=True)
refiner.unet = torch.compile(refiner.unet, mode="reduce-overhead", fullgraph=True)

如果你使用的 torch 版本 < 2.0,记得要用 xFormers 来提供 flash attention

base.enable_xformers_memory_efficient_attention()
refiner.enable_xformers_memory_efficient_attention()

Other resources

如果你想要研究下 SDXL 中 UNet2DConditionModel 的最小版本,可参考minSDXL 。

更新时间 2024-03-17