作者: Hongyu Chiu, fchollet, lukewood, divamgupta
创建日期 2024/10/09
最后修改 2024/10/24
描述: 使用 KerasHub 的 Stable Diffusion 3 模型生成图像。
Stable Diffusion 3 是一个强大的开源潜在扩散模型 (LDM),旨在根据文本提示生成高质量的新颖图像。它由 Stability AI 发布,在 10 亿张图像上进行了预训练,并在 3300 万张高质量美学和偏好图像上进行了微调,与之前版本的 Stable Diffusion 模型相比,性能得到了极大提升。
在本指南中,我们将探讨 KerasHub 对 Stable Diffusion 3 Medium 的实现,包括文本到图像、图像到图像和图像修复任务。
开始之前,我们先安装一些依赖并获取用于演示的图像
!pip install -Uq keras
!pip install -Uq git+https://github.com/keras-team/keras-hub.git
!wget -O mountain_dog.png https://raw.githubusercontent.com/keras-team/keras-io/master/guides/img/stable_diffusion_3_in_keras_hub/mountain_dog.png
!wget -O mountain_dog_mask.png https://raw.githubusercontent.com/keras-team/keras-io/master/guides/img/stable_diffusion_3_in_keras_hub/mountain_dog_mask.png
import os
os.environ["KERAS_BACKEND"] = "jax"
import time
import keras
import keras_hub
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image
在深入了解潜在扩散模型如何工作之前,我们先使用 KerasHub 的 API 生成一些图像。
为了避免为不同任务重复初始化变量,我们将使用 KerasHub 的 from_preset
工厂方法实例化并加载预训练的 backbone
和 preprocessor
。如果您只想一次执行一个任务,可以使用更简单的 API,例如:
text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
"stable_diffusion_3_medium", dtype="float16"
)
这将自动为您加载和配置预训练的 backbone
和 preprocessor
。
请注意,在本指南中,我们将使用 image_shape=(512, 512, 3)
以加快图像生成速度。对于更高质量的输出,建议使用默认尺寸 1024
。由于整个 backbone
约有 30 亿参数,这可能难以放入消费级 GPU 中,因此我们将 dtype="float16"
设置为 `float16` 以减少 GPU 内存占用——官方发布的权重也是 `float16` 格式的。
另外值得注意的是,预设 `"stable_diffusion_3_medium"` 不包含 T5XXL 文本编码器,因为它需要更多的 GPU 内存。在大多数情况下,性能下降可以忽略不计。包含 T5XXL 的权重将很快在 KerasHub 上提供。
def display_generated_images(images):
"""Helper function to display the images from the inputs.
This function accepts the following input formats:
- 3D numpy array.
- 4D numpy array: concatenated horizontally.
- List of 3D numpy arrays: concatenated horizontally.
"""
display_image = None
if isinstance(images, np.ndarray):
if images.ndim == 3:
display_image = Image.fromarray(images)
elif images.ndim == 4:
concated_images = np.concatenate(list(images), axis=1)
display_image = Image.fromarray(concated_images)
elif isinstance(images, list):
concated_images = np.concatenate(images, axis=1)
display_image = Image.fromarray(concated_images)
if display_image is None:
raise ValueError("Unsupported input format.")
plt.figure(figsize=(10, 10))
plt.axis("off")
plt.imshow(display_image)
plt.show()
plt.close()
backbone = keras_hub.models.StableDiffusion3Backbone.from_preset(
"stable_diffusion_3_medium", image_shape=(512, 512, 3), dtype="float16"
)
preprocessor = keras_hub.models.StableDiffusion3TextToImagePreprocessor.from_preset(
"stable_diffusion_3_medium"
)
text_to_image = keras_hub.models.StableDiffusion3TextToImage(backbone, preprocessor)
接下来,我们给出提示词
prompt = "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k"
# When using JAX or TensorFlow backends, you might experience a significant
# compilation time during the first `generate()` call. The subsequent
# `generate()` call speedup highlights the power of JIT compilation and caching
# in frameworks like JAX and TensorFlow, making them well-suited for
# high-performance deep learning tasks like image generation.
generated_image = text_to_image.generate(prompt)
display_generated_images(generated_image)
非常令人印象深刻!但这如何工作的呢?
让我们深入探讨“潜在扩散模型”的含义。
考虑“超分辨率”的概念,深度学习模型会“去噪”输入图像,将其转换为更高分辨率的版本。模型利用其训练数据分布来“幻化”出给定输入最可能出现的视觉细节。要了解有关超分辨率的更多信息,您可以查看以下 Keras.io 教程:
当我们把这个想法推向极致时,我们可能会开始问——如果我们只是在一个纯粹的噪声上运行这样的模型会怎样?模型会“去噪”,并开始“幻化”出一张全新的图像。通过重复这个过程多次,我们可以将一小块噪声变成越来越清晰、高分辨率的人工图像。
这就是潜在扩散的核心思想,在《高分辨率图像合成与潜在扩散模型》中提出。要深入理解扩散过程,您可以查看 Keras.io 教程 《去噪扩散隐式模型》。
为了从潜在扩散过渡到文本到图像系统,必须添加一个关键特性:使用提示关键词控制生成的视觉内容的能力。在 Stable Diffusion 3 中,使用 CLIP 和 T5XXL 模型的文本编码器获取文本嵌入,然后将其输入到扩散模型中以调节扩散过程。这种方法基于“无分类器引导”的概念,在《无分类器扩散引导》中提出。
当我们结合这些想法时,就可以对 Stable Diffusion 3 的架构有一个高层次的概览
首先,文本提示通过多个文本编码器(这些编码器是预训练且冻结的语言模型)被投影到潜在空间。接下来,文本嵌入以及随机生成的噪声块(通常来自高斯分布)被馈送到扩散模型中。扩散模型在一系列步骤中重复地对噪声块进行“去噪”(步数越多,图像越清晰、越精细——默认值为 28 步)。最后,潜在图像块通过 VAE 模型的解码器进行处理,以高分辨率渲染图像。
Stable Diffusion 3 的架构概览:
一旦我们在数十亿张图片及其标注上进行训练,这个相对简单的系统就开始看起来像魔法一样。正如费曼谈到宇宙时所说:“它并不复杂,只是量很大!”
现在我们了解了 Stable Diffusion 3 和文本到图像任务的基础知识。接下来,让我们使用 KerasHub API 进行更深入的探索。
为了使用 KerasHub 的 API 进行高效的批量处理,我们可以向模型提供一个提示词列表
generated_images = text_to_image.generate([prompt] * 3)
display_generated_images(generated_images)
num_steps
参数控制图像生成过程中使用的去噪步数。增加步数通常会产生更高质量的图像,但会增加生成时间。在 Stable Diffusion 3 中,该参数默认值为 28
。
num_steps = [10, 28, 50]
generated_images = []
for n in num_steps:
st = time.time()
generated_images.append(text_to_image.generate(prompt, num_steps=n))
print(f"Cost time (`num_steps={n}`): {time.time() - st:.2f}s")
display_generated_images(generated_images)
Cost time (`num_steps=10`): 1.35s
Cost time (`num_steps=28`): 3.44s
Cost time (`num_steps=50`): 6.18s
我们可以使用 `"negative_prompts"` 来引导模型避免生成特定的风格和元素。输入格式变为一个字典,包含 `"prompts"` 和 `"negative_prompts"` 两个键。
如果未提供 `"negative_prompts"`,则会被解释为无条件提示,默认值为 ""
。
generated_images = text_to_image.generate(
{
"prompts": [prompt] * 3,
"negative_prompts": ["Green color"] * 3,
}
)
display_generated_images(generated_images)
guidance_scale
影响 `"prompts"` 对图像生成的影响程度。较低的值赋予模型更多创造力,生成与提示词关联较宽松的图像。较高的值促使模型更紧密地遵循提示词。如果此值过高,生成的图像中可能会出现一些伪影。在 Stable Diffusion 3 中,该参数默认值为 7.0
。
generated_images = [
text_to_image.generate(prompt, guidance_scale=2.5),
text_to_image.generate(prompt, guidance_scale=7.0),
text_to_image.generate(prompt, guidance_scale=10.5),
]
display_generated_images(generated_images)
请注意,negative_prompts
和 guidance_scale
是相关的。实现中的公式可以表示如下:predicted_noise = negative_noise + guidance_scale * (positive_noise - negative_noise)
。
可以将参考图像用作扩散过程的起点。这需要在管道中额外增加一个模块:来自 VAE 模型的编码器。
参考图像由 VAE 编码器编码到潜在空间中,然后添加噪声。随后的去噪步骤与文本到图像任务的程序相同。
输入格式变为一个字典,包含 `"images"`、`"prompts"` 键,以及可选的 `"negative_prompts"` 键。
image_to_image = keras_hub.models.StableDiffusion3ImageToImage(backbone, preprocessor)
image = Image.open("mountain_dog.png").convert("RGB")
image = image.resize((512, 512))
width, height = image.size
# Note that the values of the image must be in the range of [-1.0, 1.0].
rescale = keras.layers.Rescaling(scale=1 / 127.5, offset=-1.0)
image_array = rescale(np.array(image))
prompt = "dog wizard, gandalf, lord of the rings, detailed, fantasy, cute, "
prompt += "adorable, Pixar, Disney, 8k"
generated_image = image_to_image.generate(
{
"images": image_array,
"prompts": prompt,
}
)
display_generated_images(
[
np.array(image),
generated_image,
]
)
如您所见,基于参考图像和提示词生成了新的图像。
strength
参数在决定生成的图像与参考图像的相似程度上起着关键作用。该值范围在 [0.0, 1.0]
之间,在 Stable Diffusion 3 中默认为 0.8
。
strength
值越高,模型越有“创造力”来生成与参考图像不同的图像。当值为 1.0
时,参考图像完全被忽略,任务变成纯粹的文本到图像生成。
strength
值越低,生成的图像与参考图像越相似。
generated_images = [
image_to_image.generate(
{
"images": image_array,
"prompts": prompt,
},
strength=0.7,
),
image_to_image.generate(
{
"images": image_array,
"prompts": prompt,
},
strength=0.8,
),
image_to_image.generate(
{
"images": image_array,
"prompts": prompt,
},
strength=0.9,
),
]
display_generated_images(generated_images)
在图像到图像任务的基础上,我们还可以使用蒙版来控制生成区域。这个过程称为图像修复,用于替换或编辑图像的特定区域。
图像修复依赖于蒙版来确定要修改的图像区域。要进行修复的区域由白色像素 (True
) 表示,而要保留的区域由黑色像素 (False
) 表示。
对于图像修复,输入是一个字典,包含 `"images"`、`"masks"`、`"prompts"` 键,以及可选的 `"negative_prompts"` 键。
inpaint = keras_hub.models.StableDiffusion3Inpaint(backbone, preprocessor)
image = Image.open("mountain_dog.png").convert("RGB")
image = image.resize((512, 512))
image_array = rescale(np.array(image))
# Note that the mask values are of boolean dtype.
mask = Image.open("mountain_dog_mask.png").convert("L")
mask = mask.resize((512, 512))
mask_array = np.array(mask).astype("bool")
prompt = "a black cat with glowing eyes, cute, adorable, disney, pixar, highly "
prompt += "detailed, 8k"
generated_image = inpaint.generate(
{
"images": image_array,
"masks": mask_array,
"prompts": prompt,
}
)
display_generated_images(
[
np.array(image),
np.array(mask.convert("RGB")),
generated_image,
]
)
太棒了!狗狗被一只可爱的黑猫取代了,但与图像到图像不同的是,背景得到了保留。
请注意,图像修复任务也包含 strength
参数来控制图像生成,在 Stable Diffusion 3 中默认值为 0.6
。
KerasHub 的 StableDiffusion3
支持多种应用,并在 Keras 3 的帮助下,可以在 TensorFlow、JAX 和 PyTorch 上运行模型!