Diffusers documentation
CosmosTransformer3DModel
CosmosTransformer3DModel
NVIDIAによるCosmos World Foundation Model Platform for Physical AIでは、3Dの動画のようなデータ向けのDiffusion Transformerモデルが紹介されました。
モデルは、以下のコードスニペットで読み込むことができます。
from diffusers import CosmosTransformer3DModel
transformer = CosmosTransformer3DModel.from_pretrained("nvidia/Cosmos-1.0-Diffusion-7B-Text2World", subfolder="transformer", dtype=torch.bfloat16)CosmosTransformer3DModel
class diffusers.CosmosTransformer3DModel
< source >( in_channels: int = 16out_channels: int = 16num_attention_heads: int = 32attention_head_dim: int = 128num_layers: int = 28mlp_ratio: float = 4.0text_embed_dim: int = 1024adaln_lora_dim: int = 256max_size: tuple = (128, 240, 240)patch_size: tuple = (1, 2, 2)rope_scale: tuple = (2.0, 1.0, 1.0)concat_padding_mask: bool = Trueextra_pos_embed_type: str | None = 'learnable'use_crossattn_projection: bool = Falsecrossattn_proj_in_channels: int = 1024encoder_hidden_states_channels: int = 1024controlnet_block_every_n: int | None = Noneimg_context_dim_in: int | None = Noneimg_context_num_tokens: int = 256img_context_dim_out: int = 2048 )
パラメータ
- in_channels (
int, デフォルトは16) — 入力におけるチャネル数。 - out_channels (
int、デフォルトは16) — 出力のチャネル数です。 - num_attention_heads (
int、デフォルトは32) — マルチヘッドアテンションで使用するヘッドの数。 - attention_head_dim (
int、デフォルトは128) — 各アテンションヘッドのチャネル数です。 - num_layers (
int、デフォルトは28) — 使用する transformer ブロックの層数です。 - mlp_ratio (
float、デフォルト値は4.0) — フィードフォワードネットワークにおける隠れ層サイズと入力サイズの比率です。 - text_embed_dim (
int、デフォルトは4096) — テキストエンコーダーからのテキスト埋め込みの入力次元数。 - adaln_lora_dim (
int、デフォルトは256) — Adaptive LayerNorm LoRA層の隠れ次元です。 - max_size (
tuple[int, int, int]、デフォルトは(128, 240, 240)) — 入力潜在テンソルの時間軸、高さ、幅の各次元における最大サイズです。 - patch_size (
tuple[int, int, int]、デフォルトは(1, 2, 2)) — 時間、高さ、幅の各次元で、入力潜在テンソルをパッチ化するために使用するパッチサイズ。 - rope_scale (
tuple[float, float, float]、デフォルトは(2.0, 1.0, 1.0)) — 時間、高さ、幅の各次元で RoPE に使用するスケーリング係数です。 - concat_padding_mask (
bool、デフォルトはTrue) — 入力潜在テンソルにパディングマスクを連結するかどうかを指定します。 - extra_pos_embed_type (
str、任意、デフォルトはlearnable) — 使用する追加の位置埋め込みの種類です。Noneまたはlearnableのいずれかです。 - controlnet_block_every_n (
int, 任意) — 制御残差を適用すべき Transformer ブロック間の間隔（例:7を指定すると、7 個目のブロックごとに注入されます）。Cosmos Transfer2.5 では必須です。 - img_context_dim_in (
int, 任意) — 入力画像コンテキスト特徴ベクトルの次元数です。つまり、[B, N, D] の D に相当します。 - img_context_num_tokens (
int) — 画像コンテキスト特徴ベクトルに含まれるトークン数です。つまり、[B, N, D] の N に相当します。img_context_dim_inが指定されていない場合、このパラメータは無視されます。 - img_context_dim_out (
int) — 画像コンテキスト投影層の出力次元。img_context_dim_inが指定されていない場合、このパラメータは無視されます。
Cosmosで使用される、動画のようなデータ向けのTransformerモデル。
forward
< source >( hidden_states: Tensortimestep: Tensorencoder_hidden_states: Tensorblock_controlnet_hidden_states: list[torch.Tensor] | None = Noneattention_mask: typing.Optional[torch.Tensor] = Nonefps: int | None = Nonecondition_mask: typing.Optional[torch.Tensor] = Nonepadding_mask: typing.Optional[torch.Tensor] = Nonereturn_dict: bool = True )
パラメータ
- hidden_states (
torch.Tensor、形状(batch_size, num_channels, num_frames, height, width)) — 入力となるhidden_states。 - timestep (
torch.LongTensor) — デノイジングステップを示すために使用されます。 - encoder_hidden_states (
torch.Tensor、形状(batch_size, sequence_len, embed_dims)) — 使用する条件埋め込み（プロンプトなどの入力条件から計算される埋め込み）。 - block_controlnet_hidden_states (
listoftorch.Tensor, 任意) — 指定された場合、transformer ブロックの残差に追加されるテンソルのリストです。 - attention_mask (
torch.Tensor, 任意) — アテンション時にencoder_hidden_statesに適用されるマスク。 - fps (
int, 任意) — 回転位置埋め込みを計算するために使用される入力動画の毎秒フレーム数。 - condition_mask (
torch.Tensor, 任意) —hidden_statesに連結されるマスクチャネル。条件付け領域を示します。 - padding_mask (
torch.Tensor, 任意) —concat_padding_maskが有効な場合、hidden_statesに連結されるパディングマスク。 - return_dict (
bool、任意、デフォルトはTrue) — プレーンなタプルではなく~models.transformer_2d.Transformer2DModelOutputを返すかどうか。
CosmosTransformer3DModelのforwardメソッドです。
Transformer2DModelOutput
クラス diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
パラメータ
- sample (
torch.Tensor、形状(batch_size, num_channels, height, width)、または Transformer2DModel が離散型の場合(batch size, num_vector_embeds - 1, num_latent_pixels)) —encoder_hidden_states入力を条件とする隠れ状態の出力です。離散型の場合、ノイズ除去されていない潜在ピクセルの確率分布を返します。
Transformer2DModelの出力です。