Diffusers documentation
OmniGenTransformer2DModel
OmniGenTransformer2DModel
OmniGenの画像生成に対応するマルチモーダル指示を受け取るTransformerモデルです。
論文の抄録は以下の通りです：
大規模言語モデル（LLM）の登場により、言語生成タスクが統合され、人間と機械のインタラクションは革命を起こしました。しかし、画像生成の分野では、単一のフレームワーク内で多様なタスクを処理できる統合モデルは、まだほとんど探索されていません。本稿では、統合画像生成のための新しい拡散モデルである OmniGen を紹介します。OmniGen は、以下の特性を備えています。1) 統合性: OmniGen はテキストから画像を生成する能力を示すだけでなく、画像編集、対象物駆動型生成、視覚的条件付き生成などの多様な下流タスクを本質的にサポートします。2) 簡素さ: OmniGen のアーキテクチャは高度に簡素化されており、追加のプラグインは不要です。さらに、既存の拡散モデルと比較して、よりユーザーフレンドリーであり、追加の中間ステップを必要とせずに、指示を通じて複雑なタスクをエンドツーエンドで完了でき、画像生成のワークフローを大幅に簡素化します。3) 知識の転移: 統合形式での学習の恩恵により、OmniGen は異なるタスク間で知識を効果的に転移し、未見のタスクやドメインを管理し、新しい能力を発揮します。また、モデルの推論能力と、思考連鎖（chain-of-thought）メカニズムの潜在的な応用についても検討します。本作業は、汎用画像生成モデルへの最初の試みであり、将来の発展を促進するために、リソースを 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
class 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) — パッチ埋め込み層で使用する空間パッチのサイズです。 - 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) — パディングトークンの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
< source >( hidden_states: Tensortimestep: typing.Union[int, float, torch.FloatTensor]input_ids: Tensorinput_img_latents: listinput_image_sizes: dictattention_mask: Tensorposition_ids: Tensorreturn_dict: bool = True ) → ~models.transformer_2d.Transformer2DModelOutput or tuple
パラメータ
- hidden_states (
torch.Tensor、形状(batch_size, in_channels, height, width)) — 入力となるhidden_states。 - timestep (
torch.LongTensor) — デノイジングステップを示すために使用されます。 - input_ids (
torch.Tensor) — 条件付けとして使用されるマルチモーダルなテキストトークンID。 - input_img_latents (
listoftorch.Tensor) — 条件付けに使用される入力画像の潜変数(latents)のリスト。 - input_image_sizes (
dictofinttolistofint) — サンプルのインデックスから、入力画像の埋め込みを条件付けシーケンス内のどの位置に配置すべきかを指定するマッピングです。 - 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のforwardメソッドです。