StableDiffusion3Backbone 模型

[源代码]

StableDiffusion3Backbone 类

keras_hub.models.StableDiffusion3Backbone(
    mmdit_patch_size,
    mmdit_hidden_dim,
    mmdit_num_layers,
    mmdit_num_heads,
    mmdit_position_size,
    mmdit_qk_norm,
    mmdit_dual_attention_indices,
    vae,
    clip_l,
    clip_g,
    t5=None,
    latent_channels=16,
    output_channels=3,
    num_train_timesteps=1000,
    shift=3.0,
    image_shape=(1024, 1024, 3),
    data_format=None,
    dtype=None,
    **kwargs
)

带有超参数的 Stable Diffusion 3 核心网络。

该骨干网络导入 CLIP 和 T5 模型作为文本编码器，并为 Stable Diffusion 3 模型实现基础的 MMDiT 和 VAE 网络。

默认构造函数提供一个完全可定制、随机初始化的、具有任何超参数的 MMDiT 和 VAE 模型。要加载预设的架构和权重，请使用 from_preset 构造函数。

参数

• mmdit_patch_size: 整数。MMDiT 中输入图像每个方形补丁的大小。
• mmdit_hidden_dim: 整数。MMDiT 中每个 Transformer 层末尾的 Transformer 隐藏状态的大小。
• mmdit_num_layers: 整数。MMDiT 中 Transformer 层的数量。
• mmdit_num_heads: 整数。MMDiT 中每个 Transformer 的注意力头数量。
• mmdit_position_size: 整数。MMDiT 中位置嵌入的高度和宽度的尺寸。
• mmdit_qk_norm: 可选字符串。是否对 MMDiT 中每个 Transformer 的查询和键张量进行归一化。可用选项为 None 和 "rms_norm"。通常，对于 3.0 版本，此项设置为 None；对于 3.5 版本，设置为 "rms_norm"。
• mmdit_dual_attention_indices: 可选元组。指定作为双注意力块的块的索引。通常，这用于 3.5 版本。默认为 None。
• vae: 用于在像素空间和潜在空间之间进行转换的 VAE。
• clip_l: 用于编码输入的 CLIP 文本编码器。
• clip_g: 用于编码输入的 CLIP 文本编码器。
• t5: 可选，用于编码输入的 T5 文本编码器。
• latent_channels: 整数。潜在空间中的通道数。默认为 16。
• output_channels: 整数。输出中的通道数。默认为 3。
• num_train_timesteps: 整数。训练模型的扩散步数。默认为 1000。
• shift: 浮点数。时间步调度的时间偏移值。默认为 3.0。
• image_shape: 元组。不含批量大小的输入形状。默认为 (1024, 1024, 3)。
• data_format: None 或字符串。如果指定，则为 "channels_last" 或 "channels_first"。输入中维度的顺序。"channels_last" 对应于形状为 (batch_size, height, width, channels) 的输入，而 "channels_first" 对应于形状为 (batch_size, channels, height, width) 的输入。它默认为您 Keras 配置文件 ~/.keras/keras.json 中的 image_data_format 值。如果您从未设置过它，那么它将是 "channels_last"。
• dtype: 字符串或 keras.mixed_precision.DTypePolicy。用于模型计算和权重的 dtype。请注意，无论 dtype 如何，某些计算（如 softmax 和层归一化）将始终以 float32 精度完成。

示例

# Pretrained Stable Diffusion 3 model.
model = keras_hub.models.StableDiffusion3Backbone.from_preset(
    "stable_diffusion_3_medium"
)

# Randomly initialized Stable Diffusion 3 model with custom config.
vae = keras_hub.models.VAEBackbone(...)
clip_l = keras_hub.models.CLIPTextEncoder(...)
clip_g = keras_hub.models.CLIPTextEncoder(...)
model = keras_hub.models.StableDiffusion3Backbone(
    mmdit_patch_size=2,
    mmdit_num_heads=4,
    mmdit_hidden_dim=256,
    mmdit_depth=4,
    mmdit_position_size=192,
    mmdit_qk_norm=None,
    mmdit_dual_attention_indices=None,
    vae=vae,
    clip_l=clip_l,
    clip_g=clip_g,
)

[源代码]

from_preset 方法

StableDiffusion3Backbone.from_preset(preset, load_weights=True, **kwargs)

从模型预设实例化一个 keras_hub.models.Backbone。

预设是一个包含配置、权重和其他文件资源的目录，用于保存和加载预训练模型。preset 可以作为以下之一传递：

1. 一个内置的预设标识符，如 'bert_base_en'
2. 一个 Kaggle Models 句柄，如 'kaggle://user/bert/keras/bert_base_en'
3. 一个 Hugging Face 句柄，如 'hf://user/bert_base_en'
4. 一个本地预设目录的路径，如 './bert_base_en'

此构造函数可以通过以下两种方式之一调用。可以从基类调用，如 keras_hub.models.Backbone.from_preset()，也可以从模型类调用，如 keras_hub.models.GemmaBackbone.from_preset()。如果从基类调用，返回对象的子类将从预设目录中的配置中推断出来。

对于任何 Backbone 子类，您可以运行 cls.presets.keys() 来列出该类上所有可用的内置预设。

参数

• preset：字符串。一个内置预设标识符、一个 Kaggle Models 句柄、一个 Hugging Face 句柄或一个本地目录的路径。
• load_weights：布尔值。如果为 `True`，权重将被加载到模型架构中。如果为 `False`，权重将被随机初始化。

示例

# Load a Gemma backbone with pre-trained weights.
model = keras_hub.models.Backbone.from_preset(
    "gemma_2b_en",
)

# Load a Bert backbone with a pre-trained config and random weights.
model = keras_hub.models.Backbone.from_preset(
    "bert_base_en",
    load_weights=False,
)