Hugging Face's logo

Diffusers documentation

OmniGenTransformer2DModel

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.32.2).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

OmniGenTransformer2DModel

A Transformer model that accepts multimodal instructions to generate images for OmniGen.

The abstract from the paper is:

The emergence of Large Language Models (LLMs) has unified language generation tasks and revolutionized human-machine interaction. However, in the realm of image generation, a unified model capable of handling various tasks within a single framework remains largely unexplored. In this work, we introduce OmniGen, a new diffusion model for unified image generation. OmniGen is characterized by the following features: 1) Unification: OmniGen not only demonstrates text-to-image generation capabilities but also inherently supports various downstream tasks, such as image editing, subject-driven generation, and visual conditional generation. 2) Simplicity: The architecture of OmniGen is highly simplified, eliminating the need for additional plugins. Moreover, compared to existing diffusion models, it is more user-friendly and can complete complex tasks end-to-end through instructions without the need for extra intermediate steps, greatly simplifying the image generation workflow. 3) Knowledge Transfer: Benefit from learning in a unified format, OmniGen effectively transfers knowledge across different tasks, manages unseen tasks and domains, and exhibits novel capabilities. We also explore the model’s reasoning capabilities and potential applications of the chain-of-thought mechanism. This work represents the first attempt at a general-purpose image generation model, and we will release our resources at https://github.com/VectorSpaceLab/OmniGen to foster future advancements.

import torch
from diffusers import OmniGenTransformer2DModel

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

OmniGenTransformer2DModel

class 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: typing.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' )

Parameters

  • in_channels (int, defaults to 4) — The number of channels in the input.
  • patch_size (int, defaults to 2) — The size of the spatial patches to use in the patch embedding layer.
  • hidden_size (int, defaults to 3072) — The dimensionality of the hidden layers in the model.
  • rms_norm_eps (float, defaults to 1e-5) — Eps for RMSNorm layer.
  • num_attention_heads (int, defaults to 32) — The number of heads to use for multi-head attention.
  • num_key_value_heads (int, defaults to 32) — The number of heads to use for keys and values in multi-head attention.
  • intermediate_size (int, defaults to 8192) — Dimension of the hidden layer in FeedForward layers.
  • num_layers (int, default to 32) — The number of layers of transformer blocks to use.
  • pad_token_id (int, default to 32000) — The id of the padding token.
  • vocab_size (int, default to 32064) — The size of the vocabulary of the embedding vocabulary.
  • rope_base (int, default to 10000) — The default theta value to use when creating RoPE.
  • rope_scaling (Dict, optional) — The scaling factors for the RoPE. Must contain short_factor and long_factor.
  • pos_embed_max_size (int, default to 192) — The maximum size of the positional embeddings.
  • time_step_dim (int, default to 256) — Output dimension of timestep embeddings.
  • flip_sin_to_cos (bool, default to True) — Whether to flip the sin and cos in the positional embeddings when preparing timestep embeddings.
  • downscale_freq_shift (int, default to 0) — The frequency shift to use when downscaling the timestep embeddings.
  • timestep_activation_fn (str, default to silu) — The activation function to use for the timestep embeddings.

The Transformer model introduced in OmniGen (https://arxiv.org/pdf/2409.11340).

< > Update on GitHub

←MochiTransformer3DModel PixArtTransformer2DModel→