Diffusers documentation
MiniMaxH3ImageTransformer2DModel
MiniMaxH3ImageTransformer2DModel
一种用于图像生成和编辑的 Diffusion Transformer 模型，由 MiniMax 在 MiniMax-H3-Image 中提出。
它是去除了音频流的MiniMaxH3Transformer3DModel：单个块堆栈在一个打包的一维序列上运行，该序列包含文本条件、参考图像的各行以及生成图像的各行。构建打包布局是调用方的职责，因此 forward 签名将布局与潜变量分开接收：(t, h, w) 位置网格、每行的模态标签、每行的时间步索引，以及用于寻址图像行和文本行的两个索引张量。MiniMaxH3ImageBlocks 负责构建所有这些内容。
该检查点通过并行解码蒸馏（PDD）蒸馏为六个采样步骤，每个步骤保留一个输出头。因此，forward 会接收 pdd_step_index（即当前正在运行的步骤），并使用该步骤对应的头进行投影。
import torch
from diffusers import MiniMaxH3ImageTransformer2DModel
transformer = MiniMaxH3ImageTransformer2DModel.from_pretrained(
"MiniMaxAI/MiniMax-H3-Image", subfolder="transformer_diffusers", dtype=torch.bfloat16
).to("cuda") # or "mps", "xpu", "cpu"检查点采用混合精度：输入补丁投影、时间步 MLP 和输出头为 float32，而块堆栈为 bfloat16。from_pretrained 通过 _keep_in_fp32_modules 保持该布局，因此请传入 dtype=torch.bfloat16，让它自行放置 float32 模块，而不是之后再使用 .to(torch.bfloat16) 来转换模型。
MiniMaxH3ImageTransformer2DModel
类 diffusers.MiniMaxH3ImageTransformer2DModel
< source >( num_attention_heads: int = 56attention_head_dim: int = 128hidden_size: int = 5376num_layers: int = 50num_refiner_layers: int = 2ffn_dim: int = 14336in_channels: int = 24patch_size: tuple[int, int, int] = (1, 2, 2)text_dim: int = 5120freq_dim: int = 256time_embed_hidden_dim: int = 5376time_embed_dim: int = 2688rope_freq_dim: int = 16rope_theta: float = 10000.0norm_eps: float = 1e-05qk_norm_eps: float = 1e-05final_norm_eps: float = 1e-05num_pdd_heads: int = 6 )
参数
- num_attention_heads (
int，默认为56) — 多头注意力机制中使用的注意力头数量。 - attention_head_dim (
int，默认为128) — 每个注意力头中的通道数。请注意，num_attention_heads * attention_head_dim的值 大于hidden_size。 - hidden_size (
int，默认为5376) — 打包序列的通道数（即残差流）。 - num_layers (
int，默认为50) — Transformer 块的数量。 - num_refiner_layers (
int，默认为2) — 应用于投影文本流的 token refiner 块的数量。 - ffn_dim (
int，默认为14336) — SwiGLU 前馈层的内部维度。 - in_channels (
int，默认为24) — 图像潜变量的通道数。 - patch_size (
tuple[int, int, int]，默认为(1, 2, 2)) — 用于将潜在变量打包成行的(t, h, w)补丁。 - text_dim (
int，默认为5120) — 文本编码器生成的文本条件信号的通道数。 - freq_dim (
int，默认为256) — 正弦时间步嵌入的维度。时间步在[0, 1]范围内未经缩放地被使用。 - time_embed_hidden_dim (
int，默认为5376) — 时间步 MLP 的内部维度。 - time_embed_dim (
int，默认为2688) — 时间步 MLP 的输出维度，即每个 AdaLN 投影的输入。 - rope_freq_dim (
int，默认为16) — 每个轴上的旋转频率数量。(t, h, w)轴共享一个长度为该值的inv_freq缓冲区，并且attention_head_dim通道中的2 * 3 * rope_freq_dim个通道会被旋转。 - rope_theta (
float，默认为10000.0) — 用于计算rope.inv_freq缓冲区的旋转频率调度的底数。 - norm_eps (
float，默认为1e-5) — 注意力前和前馈层前的归一化 epsilon。 - qk_norm_eps (
float，默认为1e-5) — 每个头 query/key 范数的 epsilon。 - final_norm_eps (
float，默认为1e-5) — token refiner 输出归一化以及norm_out的 epsilon。 - num_pdd_heads (
int，默认为6) — 检查点被蒸馏到的采样步数，每个步数都有自己独立的输出头。
用于图像生成和编辑的 Transformer 模型，在 MiniMax-H3-Image 中引入。
这是去除了音频流的 MiniMax-H3 变换器：单个块堆栈运行在一个打包的一维序列上，该序列包含文本条件、参考图像的各行以及生成图像的各行。注意力机制是该序列上的完整自注意力，没有交叉注意力。检查点通过并行解码蒸馏（PDD）蒸馏为固定数量的采样步骤，每个步骤保留一个输出头，因此forward 接收其正在运行的步骤的索引。
调用方负责构建打包布局：对潜变量进行分块（patchify）、对行进行排序，并生成 (t, h, w) 位置网格、每行的模态标签以及每行的时间步索引。序列不包含填充，因此注意力机制在单个文档上无掩码运行，所有注意力后端均可用。
批次轴是一个纯粹的复制轴：结构参数（timestep、timestep_indices、token_tags、position_ids 以及两个索引张量）描述了一种打包布局，该布局由批次中的每个样本共享。
forward
< source >( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensortimestep_indices: Tensortoken_tags: Tensorposition_ids: Tensorimage_indices: Tensortext_indices: Tensorpdd_step_index: intreturn_dict: bool = True ) → Transformer2DModelOutput or tuple
参数
- hidden_states (
torch.Tensor，形状为(batch_size, num_image_tokens, in_channels * prod(patch_size))) — 分块的潜在行——参考图像的行和生成图像的行——按照它们在打包序列中出现的顺序排列，即与image_indices匹配。 - encoder_hidden_states（形状为
(batch_size, num_text_tokens, text_dim)的torch.Tensor）— 文本条件，其顺序与text_indices相匹配。 - timestep（形状为
(num_timesteps,)的torch.Tensor）— 打包序列中存在的不同时间步值，取值范围为[0, 1]且未缩放。一次前向传播服务于处于不同噪声水平的行（生成的行、参考行）。 - timestep_indices（形状为
(seq_len,)的torch.Tensor）— 对于打包序列的每一行，其在timestep中对应的时间步索引。 - token_tags（形状为
(seq_len,)的torch.Tensor）— 对于打包序列的每一行，其模态：0表示图像，1表示文本。 - position_ids（形状为
(seq_len, 3)的torch.Tensor）— 打包序列中每一行的(t, h, w)旋转坐标。 - image_indices（形状为
(num_image_tokens,)的torch.Tensor）— 图像行在打包序列中的位置。 - text_indices（形状为
(num_text_tokens,)的torch.Tensor）— 文本行在打包序列中的位置。 - pdd_step_index (
int) — 采样步骤，用于选择输出头。 - return_dict (
bool，默认为True) — 是否返回 Transformer2DModelOutput 而不是普通元组。
返回值
Transformer2DModelOutput 或 tuple
image_indices 中按行顺序排列的每个图像行的速度，形状为 (batch_size, num_image_tokens, in_channels * prod(patch_size))。