Diffusers documentation
CosmosTransformer3DModel
CosmosTransformer3DModel
Модель Diffusion Transformer для 3D-видеоподобных данных была представлена в Cosmos World Foundation Model Platform for Physical AI от NVIDIA.
Модель можно загрузить с помощью следующего фрагмента кода.
from diffusers import CosmosTransformer3DModel
transformer = CosmosTransformer3DModel.from_pretrained("nvidia/Cosmos-1.0-Diffusion-7B-Text2World", subfolder="transformer", dtype=torch.bfloat16)CosmosTransformer3DModel
класс 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[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)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) — Количество блоков трансформера, которые будут использоваться. - 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, необязательный) — Интервал между блоками трансформера, которые должны получать управляющие остаточные сигналы (например,7для инъекции после каждого седьмого блока). Обязательный параметр для Cosmos Transfer2.5. - img_context_dim_in (
int, необязательный) — Размерность вектора признаков контекста входного изображения, то есть это D в [B, N, D]. - img_context_num_tokens (
int) — Количество токенов в векторе признаков контекста изображения, то есть это N в [B, N, D]. Еслиimg_context_dim_inне задан, этот параметр игнорируется. - img_context_dim_out (
int) — Выходное измерение слоя проекции контекста изображения. Еслиimg_context_dim_inне задано, этот параметр игнорируется.
Модель Transformer для видеоподобных данных, используемая в Cosmos.
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, необязательный) — Список тензоров, которые при указании добавляются к остаточным связям блоков трансформера. - attention_mask (
torch.Tensor, необязательный) — Маска, применяемая кencoder_hidden_statesво время внимания. - fps (
int, необязательный) — Частота кадров входного видео, используемая для вычисления вращательных позиционных эмбеддингов. - condition_mask (
torch.Tensor, необязательный) — Канал маски, добавляемый кhidden_statesдля указания области условного контекста. - padding_mask (
torch.Tensor, необязательный) — Маска отступов, добавляемая кhidden_states, если включёнconcat_padding_mask. - return_dict (
bool, необязательный, по умолчаниюTrue) — Определяет, следует ли возвращать~models.transformer_2d.Transformer2DModelOutputвместо обычного кортежа.
Метод forward модели CosmosTransformer3DModel.
Transformer2DModelOutput
класс diffusers.models.modeling_outputs.Transformer2DModelOutput
< source >( sample: torch.Tensor )
Параметры
- sample (
torch.Tensorформы(batch_size, num_channels, height, width)или(batch size, num_vector_embeds - 1, num_latent_pixels), если Transformer2DModel дискретный) — Выходные скрытые состояния, обусловленные входомencoder_hidden_states. Если модель дискретная, возвращает распределения вероятностей для латентных пикселей без шума.
Вывод Transformer2DModel.