Diffusers documentation

CosmosTransformer3DModel

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

and get access to the augmented documentation experience

to get started

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

< >

( 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) — Количество блоков трансформера, которые будут использоваться.
  • 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

< >

( 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 (list of torch.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

< >

( 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.

Update on GitHub