Diffusers documentation
OmniGenTransformer2DModel
OmniGenTransformer2DModel
一个能够接收多模态指令并为 OmniGen 生成图像的 Transformer 模型。
该论文的摘要如下：
大语言模型（LLM）的出现统一了语言生成任务，并彻底改变了人机交互方式。然而，在图像生成领域，能够在单一框架内处理多种任务的统一模型仍鲜有探索。在本工作中，我们提出了 OmniGen，一种用于统一图像生成的新型扩散模型。OmniGen 具有以下特点：1) 统一性：OmniGen 不仅具备文本到图像的生成能力，还天然支持多种下游任务，例如图像编辑、主体驱动生成和视觉条件生成。2) 简洁性：OmniGen 的架构高度简化，无需额外的插件。此外，与现有的扩散模型相比，它更易于使用，能够通过指令端到端地完成复杂任务，无需额外的中间步骤，从而大大简化了图像生成工作流。3) 知识迁移：得益于统一格式下的学习，OmniGen 能够有效地在不同任务之间迁移知识，处理未见过的任务和领域，并展现出新颖的能力。我们还探索了模型的推理能力以及思维链机制的潜在应用。本工作是对通用图像生成模型的首次尝试，我们将在 https://github.com/VectorSpaceLab/OmniGen 发布我们的资源，以促进未来的发展。
import torch
from diffusers import OmniGenTransformer2DModel
transformer = OmniGenTransformer2DModel.from_pretrained("Shitao/OmniGen-v1-diffusers", subfolder="transformer", dtype=torch.bfloat16)OmniGenTransformer2DModel
类 diffusers.OmniGenTransformer2DModel
< source >( in_channels: int = 4patch_size: int = 2hidden_size: int = 3072rms_norm_eps: float = 1e-05num_attention_heads: int = 32num_key_value_heads: int = 32intermediate_size: int = 8192num_layers: int = 32pad_token_id: int = 32000vocab_size: int = 32064max_position_embeddings: int = 131072original_max_position_embeddings: int = 4096rope_base: int = 10000rope_scaling: dict = Nonepos_embed_max_size: int = 192time_step_dim: int = 256flip_sin_to_cos: bool = Truedownscale_freq_shift: int = 0timestep_activation_fn: str = 'silu' )
参数
- in_channels (
int，默认为4) — 输入数据的通道数。 - patch_size (
int，默认为2) — 用于 patch 嵌入层的空间 patch 大小。 - hidden_size (
int，默认为3072) — 模型中隐藏层的维度。 - rms_norm_eps (
float，默认为1e-5) — RMSNorm 层的 Eps。 - num_attention_heads (
int，默认为32) — 多头注意力机制中使用的注意力头数量。 - num_key_value_heads (
int，默认为32) — 多头注意力中用于键和值的头数。 - intermediate_size (
int，默认为8192) — FeedForward 层中隐藏层的维度。 - num_layers (
int，默认为32) — 要使用的 transformer 块层数。 - pad_token_id (
int，默认为32000) — 填充 token 的 id。 - vocab_size (
int，默认为32064) — 嵌入词表的大小。 - rope_base (
int，默认为10000) — 创建 RoPE 时使用的默认 theta 值。 - rope_scaling (
dict, 可选) — RoPE 的缩放因子。必须包含short_factor和long_factor。 - pos_embed_max_size (
int，默认为192) — 位置嵌入的最大尺寸。 - time_step_dim (
int，默认为256) — 时间步嵌入的输出维度。 - flip_sin_to_cos (
bool，默认为True) — 在准备时间步嵌入时，是否交换位置嵌入中的 sin 和 cos。 - downscale_freq_shift（
int，默认为0）— 在对时间步嵌入进行下采样时使用的频率偏移。 - timestep_activation_fn (
str，默认为silu) — 用于时间步嵌入的激活函数。
OmniGen 中介绍的 Transformer 模型（https://huggingface.co/papers/2409.11340）。
forward
< source >( hidden_states: Tensortimestep: typing.Union[int, float, torch.FloatTensor]input_ids: Tensorinput_img_latents: list[torch.Tensor]input_image_sizes: dict[int, list[int]]attention_mask: Tensorposition_ids: Tensorreturn_dict: bool = True ) → ~models.transformer_2d.Transformer2DModelOutput or tuple
参数
- hidden_states（形状为
(batch_size, in_channels, height, width)的torch.Tensor）— 输入的hidden_states。 - timestep (
torch.LongTensor) — 用于指示去噪步骤。 - input_ids (
torch.Tensor) — 用作条件输入的多模态文本 token id。 - input_img_latents (
listoftorch.Tensor) — 用作条件输入的图像潜变量列表。 - input_image_sizes（
dict，键为int，值为listofint） — 从样本索引到输入图像嵌入在条件序列中应放置位置的映射。 - attention_mask (
torch.Tensor) — 用于联合多模态序列的注意力掩码。 - position_ids (
torch.Tensor) — 用于计算位置嵌入的位置 ID。 - return_dict (
bool，可选，默认为True) — 是否返回~models.transformer_2d.Transformer2DModelOutput而不是普通的元组。
返回值
~models.transformer_2d.Transformer2DModelOutput 或 tuple
如果 return_dict 为 True，则返回 ~models.transformer_2d.Transformer2DModelOutput，否则返回普通的 tuple。
OmniGenTransformer2DModel 的前向方法。