#
OmniGenTransformer2DModel

A Transformer model accept multi-modal instruction to generate image from OmniGen.

##
OmniGenTransformer2DModel

###
class diffusers.OmniGenTransformer2DModel

< source >

( transformer_config: typing.Dict patch_size = 2 in_channels = 4 pos_embed_max_size: int = 192 )

Parameters

- transformer_config (`dict`) — config for transformer layers. OmniGen-v1 use Phi3 as transformer backbone
- patch_size (`int`, defaults to 2) — Patch size to turn the input data into small patches.
- in_channels (`int`, optional, defaults to 4) — The number of channels in the input.
- pos_embed_max_size (`int`, optional, defaults to 192) — The max size of pos emb.

The Transformer model introduced in OmniGen.

Reference: https://arxiv.org/pdf/2409.11340

####
forward

< source >

( hidden_states: Tensor timestep: typing.Union[int, float, torch.LongTensor] input_ids: Tensor input_img_latents: typing.List[torch.Tensor] input_image_sizes: typing.Dict[int, typing.List[int]] attention_mask: Tensor position_ids: Tensor past_key_values: DynamicCache = None offload_transformer_block: bool = False attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None return_dict: bool = True )

Parameters

- hidden_states (`torch.FloatTensor` of shape `(batch size, channel, height, width)`) — Input `hidden_states`.
- timestep (`torch.LongTensor`) — Used to indicate denoising step.
- input_ids (`torch.LongTensor`) — token ids
- input_img_latents (`torch.FloatTensor`) — encoded image latents by VAE
- input_image_sizes (`dict`) — the indices of the input_img_latents in the input_ids
- attention_mask (`torch.FloatTensor`) — mask for self-attention
- position_ids (`torch.LongTensor`) — id to represent position
- past_key_values (`transformers.cache_utils.Cache`) — previous key and value states
- offload_transformer_block (`bool`, optional, defaults to `True`) — offload transformer block to cpu
- attention_kwargs — (`dict`, optional): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in diffusers.models.attention_processor.
- return_dict (`bool`, optional, defaults to `True`) — Whether or not to return a `OmniGen2DModelOutput` instead of a plain tuple.

The OmniGenTransformer2DModel forward method.

####
get_multimodal_embeddings

< source >

( input_ids: Tensor input_img_latents: typing.List[torch.Tensor] input_image_sizes: typing.Dict )

Parameters

- input_ids — a sequence of text id
- input_img_latents — continues embedding of input images
- input_image_sizes — the index of the input image in the input_ids sequence.

get the multi-modal conditional embeddings

Returns: torch.Tensor

####
set_attn_processor

< source >

( processor: typing.Union[diffusers.models.attention_processor.AttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor2_0, diffusers.models.attention_processor.JointAttnProcessor2_0, diffusers.models.attention_processor.PAGJointAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGJointAttnProcessor2_0, diffusers.models.attention_processor.FusedJointAttnProcessor2_0, diffusers.models.attention_processor.AllegroAttnProcessor2_0, diffusers.models.attention_processor.AuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FusedAuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.CogVideoXAttnProcessor2_0, diffusers.models.attention_processor.FusedCogVideoXAttnProcessor2_0, diffusers.models.attention_processor.XFormersAttnAddedKVProcessor, diffusers.models.attention_processor.XFormersAttnProcessor, diffusers.models.attention_processor.XLAFlashAttnProcessor2_0, diffusers.models.attention_processor.AttnProcessorNPU, diffusers.models.attention_processor.AttnProcessor2_0, diffusers.models.attention_processor.MochiVaeAttnProcessor2_0, diffusers.models.attention_processor.StableAudioAttnProcessor2_0, diffusers.models.attention_processor.HunyuanAttnProcessor2_0, diffusers.models.attention_processor.FusedHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.LuminaAttnProcessor2_0, diffusers.models.attention_processor.MochiAttnProcessor2_0, diffusers.models.attention_processor.FusedAttnProcessor2_0, diffusers.models.attention_processor.CustomDiffusionXFormersAttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor2_0, diffusers.models.attention_processor.SlicedAttnProcessor, diffusers.models.attention_processor.SlicedAttnAddedKVProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor2_0, diffusers.models.attention_processor.IPAdapterXFormersAttnProcessor, diffusers.models.attention_processor.PAGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.LoRAAttnProcessor, diffusers.models.attention_processor.LoRAAttnProcessor2_0, diffusers.models.attention_processor.LoRAXFormersAttnProcessor, diffusers.models.attention_processor.LoRAAttnAddedKVProcessor, typing.Dict[str, typing.Union[diffusers.models.attention_processor.AttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor, diffusers.models.attention_processor.AttnAddedKVProcessor2_0, diffusers.models.attention_processor.JointAttnProcessor2_0, diffusers.models.attention_processor.PAGJointAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGJointAttnProcessor2_0, diffusers.models.attention_processor.FusedJointAttnProcessor2_0, diffusers.models.attention_processor.AllegroAttnProcessor2_0, diffusers.models.attention_processor.AuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FusedAuraFlowAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0, diffusers.models.attention_processor.FluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0, diffusers.models.attention_processor.FusedFluxAttnProcessor2_0_NPU, diffusers.models.attention_processor.CogVideoXAttnProcessor2_0, diffusers.models.attention_processor.FusedCogVideoXAttnProcessor2_0, diffusers.models.attention_processor.XFormersAttnAddedKVProcessor, diffusers.models.attention_processor.XFormersAttnProcessor, diffusers.models.attention_processor.XLAFlashAttnProcessor2_0, diffusers.models.attention_processor.AttnProcessorNPU, diffusers.models.attention_processor.AttnProcessor2_0, diffusers.models.attention_processor.MochiVaeAttnProcessor2_0, diffusers.models.attention_processor.StableAudioAttnProcessor2_0, diffusers.models.attention_processor.HunyuanAttnProcessor2_0, diffusers.models.attention_processor.FusedHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGHunyuanAttnProcessor2_0, diffusers.models.attention_processor.LuminaAttnProcessor2_0, diffusers.models.attention_processor.MochiAttnProcessor2_0, diffusers.models.attention_processor.FusedAttnProcessor2_0, diffusers.models.attention_processor.CustomDiffusionXFormersAttnProcessor, diffusers.models.attention_processor.CustomDiffusionAttnProcessor2_0, diffusers.models.attention_processor.SlicedAttnProcessor, diffusers.models.attention_processor.SlicedAttnAddedKVProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor, diffusers.models.attention_processor.IPAdapterAttnProcessor2_0, diffusers.models.attention_processor.IPAdapterXFormersAttnProcessor, diffusers.models.attention_processor.PAGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.PAGCFGIdentitySelfAttnProcessor2_0, diffusers.models.attention_processor.LoRAAttnProcessor, diffusers.models.attention_processor.LoRAAttnProcessor2_0, diffusers.models.attention_processor.LoRAXFormersAttnProcessor, diffusers.models.attention_processor.LoRAAttnAddedKVProcessor]]] )

Parameters

- processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`) — The instantiated processor class or a dictionary of processor classes that will be set as the processor for all `Attention` layers.

  If `processor` is a dict, the key needs to define the path to the corresponding cross attention processor. This is strongly recommended when setting trainable attention processors.

Sets the attention processor to use to compute attention.

####
unpatchify

< source >

( x h w )

x: (N, T, patch_size*2 C) imgs: (N, H, W, C)

< > Update on GitHub