KerasHub:预训练模型 / 开发者指南 / KerasHub 中的 Stable Diffusion 3!

KerasHub 中的 Stable Diffusion 3!

作者: Hongyu Chiu, fchollet, lukewood, divamgupta
创建日期 2024/10/09
最后修改 2024/10/24
描述: 使用 KerasHub 的 Stable Diffusion 3 模型生成图像。

在 Colab 中查看 GitHub 源代码


概览

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 工厂方法实例化并加载预训练的 backbonepreprocessor。如果您只想一次执行一个任务,可以使用更简单的 API,例如:

text_to_image = keras_hub.models.StableDiffusion3TextToImage.from_preset(
    "stable_diffusion_3_medium", dtype="float16"
)

这将自动为您加载和配置预训练的 backbonepreprocessor

请注意,在本指南中,我们将使用 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)

png

非常令人印象深刻!但这如何工作的呢?

让我们深入探讨“潜在扩散模型”的含义。

考虑“超分辨率”的概念,深度学习模型会“去噪”输入图像,将其转换为更高分辨率的版本。模型利用其训练数据分布来“幻化”出给定输入最可能出现的视觉细节。要了解有关超分辨率的更多信息,您可以查看以下 Keras.io 教程:

Super-resolution

当我们把这个想法推向极致时,我们可能会开始问——如果我们只是在一个纯粹的噪声上运行这样的模型会怎样?模型会“去噪”,并开始“幻化”出一张全新的图像。通过重复这个过程多次,我们可以将一小块噪声变成越来越清晰、高分辨率的人工图像。

这就是潜在扩散的核心思想,在《高分辨率图像合成与潜在扩散模型》中提出。要深入理解扩散过程,您可以查看 Keras.io 教程 《去噪扩散隐式模型》

Denoising diffusion

为了从潜在扩散过渡到文本到图像系统,必须添加一个关键特性:使用提示关键词控制生成的视觉内容的能力。在 Stable Diffusion 3 中,使用 CLIP 和 T5XXL 模型的文本编码器获取文本嵌入,然后将其输入到扩散模型中以调节扩散过程。这种方法基于“无分类器引导”的概念,在《无分类器扩散引导》中提出。

当我们结合这些想法时,就可以对 Stable Diffusion 3 的架构有一个高层次的概览

  • 文本编码器:将文本提示转换为文本嵌入。
  • 扩散模型:重复地对较小的潜在图像块进行“去噪”。
  • 解码器:将最终的潜在图像块转换为更高分辨率的图像。

首先,文本提示通过多个文本编码器(这些编码器是预训练且冻结的语言模型)被投影到潜在空间。接下来,文本嵌入以及随机生成的噪声块(通常来自高斯分布)被馈送到扩散模型中。扩散模型在一系列步骤中重复地对噪声块进行“去噪”(步数越多,图像越清晰、越精细——默认值为 28 步)。最后,潜在图像块通过 VAE 模型的解码器进行处理,以高分辨率渲染图像。

Stable Diffusion 3 的架构概览:Stable Diffusion 3 架构

一旦我们在数十亿张图片及其标注上进行训练,这个相对简单的系统就开始看起来像魔法一样。正如费曼谈到宇宙时所说:“它并不复杂,只是量很大!”


文本到图像任务

现在我们了解了 Stable Diffusion 3 和文本到图像任务的基础知识。接下来,让我们使用 KerasHub API 进行更深入的探索。

为了使用 KerasHub 的 API 进行高效的批量处理,我们可以向模型提供一个提示词列表

generated_images = text_to_image.generate([prompt] * 3)
display_generated_images(generated_images)

png

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

png

我们可以使用 `"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)

png

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)

png

请注意,negative_promptsguidance_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,
    ]
)

png

如您所见,基于参考图像和提示词生成了新的图像。

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)

png


图像修复任务

在图像到图像任务的基础上,我们还可以使用蒙版来控制生成区域。这个过程称为图像修复,用于替换或编辑图像的特定区域。

图像修复依赖于蒙版来确定要修改的图像区域。要进行修复的区域由白色像素 (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,
    ]
)

png

太棒了!狗狗被一只可爱的黑猫取代了,但与图像到图像不同的是,背景得到了保留。

请注意,图像修复任务也包含 strength 参数来控制图像生成,在 Stable Diffusion 3 中默认值为 0.6


结论

KerasHub 的 StableDiffusion3 支持多种应用,并在 Keras 3 的帮助下,可以在 TensorFlow、JAX 和 PyTorch 上运行模型!