# OmniGenTransformer2DModel

一个能够接收多模态指令并为 [OmniGen](https://github.com/VectorSpaceLab/OmniGen/) 生成图像的 Transformer 模型。

该论文的摘要如下：

*大语言模型（LLM）的出现统一了语言生成任务，并彻底改变了人机交互方式。然而，在图像生成领域，能够在单一框架内处理多种任务的统一模型仍鲜有探索。在本工作中，我们提出了 OmniGen，一种用于统一图像生成的新型扩散模型。OmniGen 具有以下特点：1) 统一性：OmniGen 不仅具备文本到图像的生成能力，还天然支持多种下游任务，例如图像编辑、主体驱动生成和视觉条件生成。2) 简洁性：OmniGen 的架构高度简化，无需额外的插件。此外，与现有的扩散模型相比，它更易于使用，能够通过指令端到端地完成复杂任务，无需额外的中间步骤，从而大大简化了图像生成工作流。3) 知识迁移：得益于统一格式的学习，OmniGen 能够有效地在不同任务之间迁移知识，处理未见过的任务和领域，并展现出新颖的能力。我们还探索了模型的推理能力以及思维链机制的潜在应用。本工作是对通用图像生成模型的首次尝试，我们将在 https://github.com/VectorSpaceLab/OmniGen 发布我们的资源，以促进未来的发展。*

```python
import torch
from diffusers import OmniGenTransformer2DModel

transformer = OmniGenTransformer2DModel.from_pretrained("Shitao/OmniGen-v1-diffusers", subfolder="transformer", dtype=torch.bfloat16)
```

## OmniGenTransformer2DModel[[diffusers.OmniGenTransformer2DModel]]

#### diffusers.OmniGenTransformer2DModel[[diffusers.OmniGenTransformer2DModel]]

```python
diffusers.OmniGenTransformer2DModel(in_channels: int = 4, patch_size: int = 2, hidden_size: int = 3072, rms_norm_eps: float = 1e-05, num_attention_heads: int = 32, num_key_value_heads: int = 32, intermediate_size: int = 8192, num_layers: int = 32, pad_token_id: int = 32000, vocab_size: int = 32064, max_position_embeddings: int = 131072, original_max_position_embeddings: int = 4096, rope_base: int = 10000, rope_scaling: dict = None, pos_embed_max_size: int = 192, time_step_dim: int = 256, flip_sin_to_cos: bool = True, downscale_freq_shift: int = 0, timestep_activation_fn: str = 'silu')
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_omnigen.py#L283)

**参数：**

in_channels (`int`，默认为 `4`)：输入中的通道数。

patch_size (`int`，默认为 `2`)：在 patch embedding 层中使用的空间 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 (https://huggingface.co/papers/2409.11340) 中引入的 Transformer 模型。

#### forward[[diffusers.OmniGenTransformer2DModel.forward]]

```python
forward(hidden_states: Tensor, timestep: typing.Union[int, float, torch.FloatTensor], input_ids: Tensor, input_img_latents: list[torch.Tensor], input_image_sizes: dict[int, list[int]], attention_mask: Tensor, position_ids: Tensor, return_dict: bool = True)
```

[源代码](https://github.com/huggingface/diffusers/blob/v0.43.0/src/diffusers/models/transformers/transformer_omnigen.py#L407)

**参数：**

hidden_states (`torch.Tensor`，形状为 `(batch_size, in_channels, height, width)`)：输入 `hidden_states`。

timestep (`torch.LongTensor`)：用于指示去噪步骤。

input_ids (`torch.Tensor`) : 用作条件输入的多模态文本 token id。

input_img_latents (`list` of `torch.Tensor`) : 用作条件输入的图像潜变量列表。

input_image_sizes (`dict`，键为 `int`，值为 `list` of `int`) : 从样本索引到条件序列中应放置输入图像嵌入的位置的映射。

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](/docs/diffusers/v0.43.0/zh/api/models/omnigen_transformer#diffusers.OmniGenTransformer2DModel) 的前向方法。