Diffusers documentation

AnyFlowFARTransformer3DModel

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

AnyFlowFARTransformer3DModel

O Transformer 3D causal (FAR) usado por AnyFlowFARPipeline — a variante FAR do AnyFlow. Consulte a página do AnyFlowFARPipeline para o paper, os autores e os checkpoints liberados. Ele estende o backbone Wan2.1 da v0.35.1 com três adições:

  1. máscara de bloco causal FAR via torch.nn.attention.flex_attention, com suporte à geração autoregressiva por blocos, conforme introduzido em FAR.
  2. Embedding de patches de frames comprimidos (far_patch_embedding) para frames de contexto (já gerados), inicializado a partir do patch_embedding de resolução completa no momento da construção, por interpolação trilinear.
  3. Embedding de fluxo-mapa com dois passos de tempo (igual ao AnyFlowTransformer3DModel) — cada chamada de forward é condicionada tanto no passo de tempo de origem t quanto no passo de tempo de destino r.

O cronograma de blocos padrão (chunk_partition) é armazenado na configuração do modelo; os checkpoints NVIDIA AnyFlow-FAR publicados usam [1, 3, 3, 3, 3, 3, 3, 2] para a configuração canônica de 81 quadros. forward aceita uma substituição de chunk_partition por chamada, de modo que o mesmo checkpoint também lida com outras configurações de num_frames sem necessidade de re-treinamento.

from diffusers import AnyFlowFARTransformer3DModel

# Causal AnyFlow checkpoint (FAR):
transformer = AnyFlowFARTransformer3DModel.from_pretrained(
    "nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers", subfolder="transformer"
)

AnyFlowFARTransformer3DModel

classe diffusers.AnyFlowFARTransformer3DModel

< >

( patch_size: typing.Tuple[int] = (1, 2, 2)compressed_patch_size: typing.Tuple[int] = (1, 4, 4)full_chunk_limit: int = 3num_attention_heads: int = 40attention_head_dim: int = 128in_channels: int = 16out_channels: int = 16text_dim: int = 4096freq_dim: int = 256ffn_dim: int = 13824num_layers: int = 40cross_attn_norm: bool = Trueeps: float = 1e-06image_dim: typing.Optional[int] = Nonerope_max_seq_len: int = 1024gate_value: float = 0.25deltatime_type: str = 'r'chunk_partition: typing.Tuple[int, ...] = (1, 3, 3, 3, 3, 3, 3, 2) )

Parâmetros

  • patch_size (Tuple[int], padrão (1, 2, 2)) — Dimensões do patch 3D para blocos em resolução completa.
  • compressed_patch_size (Tuple[int], padrão (1, 4, 4)) — Dimensões de patch maiores para os blocos (contexto) comprimidos pelo FAR.
  • full_chunk_limit (int, padrão 3) — Número máximo de blocos em resolução completa antes que blocos anteriores sejam rebaixados para o contexto FAR comprimido. Os checkpoints liberados usam 3.
  • num_attention_heads (int, padrão 40) — Número de cabeças de atenção.
  • attention_head_dim (int, padrão 128) — O número de canais em cada head.
  • in_channels (int, padrão 16) — O número de canais no latente de entrada.
  • out_channels (int, padrão 16) — O número de canais no latente de saída.
  • text_dim (int, padrão 4096) — Dimensão de entrada para embeddings de texto (UMT5).
  • freq_dim (int, padrão 256) — Dimensão para incorporações de tempo sinusoidais.
  • ffn_dim (int, padrão 13824) — Dimensão intermediária na rede feed-forward.
  • num_layers (int, padrão 40) — Número de blocos de transformer.
  • cross_attn_norm (bool, padrão True) — Ativa a normalização de atenção cruzada.
  • eps (float, padrão 1e-6) — Epsilon para as camadas de normalização.
  • image_dim (Optional[int], opcional, padrão None) — Dimensão do embedding de imagem para condicionamento I2V.
  • rope_max_seq_len (int, padrão 1024) — Comprimento máximo da sequência usado para pré-computar as frequências de posição rotativa.
  • gate_value (float, padrão 0.25) — Porta de mistura entre os embeddings de timestep de origem e timestep delta.
  • deltatime_type (str, padrão ‘r’) — Ou "r" (o delta é o timestep alvo) ou "t-r" (o delta é o intervalo absoluto).
  • chunk_partition (Tuple[int, …], padrão (1, 3, 3, 3, 3, 3, 3, 2)) — Quantidade padrão de frames por chunk usada pelo pipeline. Os checkpoints NVIDIA AnyFlow-FAR liberados têm como alvo num_frames=81 (21 frames latentes com passo temporal 4 do VAE), divididos como 1 + 3*6 + 2. Um num_frames diferente exige uma substituição correspondente de chunk_partition passada para AnyFlowFARPipeline.call() (e da mesma forma para forward).

Transformer 3D causal (FAR) para amostragem de mapas de fluxo do AnyFlow com geração autoregressiva por blocos.

Estende o backbone Wan2.1 v0.35.1 com:

  • máscara de bloco causal FAR via torch.nn.attention.flex_attention, com suporte à geração autoregressiva por blocos (FAR).
  • Embedding de patches de frames comprimidos far_patch_embedding para frames de contexto (já gerados), inicializado a partir de patch_embedding via interpolação trilinear, de modo que um modelo recém-construído já esteja em um ponto de partida razoável mesmo antes do fine-tuning com LoRA.
  • Embedding de fluxo-mapa com dois passos de tempo para amostragem em qualquer passo (igual a AnyFlowTransformer3DModel).

Use AnyFlowTransformer3DModel em vez disso para T2V bidirecional puro — essa variante ignora o mascaramento causal FAR e far_patch_embedding e é ~5–10% menor.

build_attention_mask

< >

( chunk_partition: typing.List[int]height: intwidth: inthas_clean_context: bool = Falsedevice: typing.Optional[torch.device] = Nonemode: str = 'train' ) → ~torch.nn.attention.flex_attention.BlockMask

Parâmetros

  • chunk_partition — contagem de quadros por chunk (a soma deve ser igual ao número de quadros latentes).
  • height, width — dimensões espaciais latentes.
  • has_clean_context — True quando clean_hidden_states será passado por forward (treino de V2V/I2V); apenas essa flag de presença afeta o layout da máscara.
  • device — dispositivo para o BlockMask resultante. A máscara não é movida automaticamente por device_map="auto"; construa-a no mesmo dispositivo em que as entradas do transformer ficarão.
  • mode — "train" (corresponde a _forward_train) ou "cache" (corresponde a _forward_cache). O caminho autoregressivo _forward_inference atende através do cache KV e não possui modo aqui.

Retorna

~torch.nn.attention.flex_attention.BlockMask

máscara causal abrangendo o layout FAR, preenchida até um múltiplo de 128 ao longo da dimensão da sequência (o requisito de tamanho de bloco do BlockMask).

Exceções

ValueError

  • ValueError — se mode não for nem "train" nem "cache".

Pré-construa a ~torch.nn.attention.flex_attention.BlockMask causal fora do forward.

Passe o resultado via o kwarg attention_mask de forward para tornar o transformer inteiro compatível com torch.compile(fullgraph=True). Sem uma máscara pré-construída, forward recorre à construção interna dela — esse caminho usa flex_attention.create_block_mask(_compile=False) e quebra o grafo de compilação.

forward

< >

( hidden_states: Tensortimestep: Tensorr_timestep: Tensorencoder_hidden_states: Tensorchunk_partition: typing.List[int]encoder_hidden_states_image: typing.Optional[torch.Tensor] = Noneclean_hidden_states: typing.Optional[torch.Tensor] = Noneclean_timestep: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = Nonekv_cache_flag: typing.Optional[typing.Dict[str, typing.Any]] = Noneattention_mask: typing.Optional[torch.nn.attention.flex_attention.BlockMask] = Noneattention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = Nonereturn_dict: bool = True ) → [~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] or tuple

Parâmetros

  • hidden_states (torch.Tensor) — Entrada latente de forma (B, F, C, H, W).
  • timestep (torch.Tensor) — Timestep t do fluxo-mapa de origem (mais ruidoso).
  • r_timestep (torch.Tensor) — Timestep r do fluxo-mapa alvo (mais limpo).
  • encoder_hidden_states (torch.Tensor) — Embutimentos de texto do UMT5.
  • chunk_partition (List[int]) — Contagem de quadros por chunk; o total deve corresponder ao número de quadros latentes em hidden_states.
  • encoder_hidden_states_image (torch.Tensor, opcional) — Embedding de imagem I2V; concatenado antes dos tokens de texto quando fornecido.
  • clean_hidden_states (torch.Tensor, opcional) — Quadros de condicionamento limpos (sem ruído) usados pelo rollout de treinamento.
  • clean_timestep (torch.Tensor, opcional) — Timesteps para os quadros de condicionamento limpo no rollout de treinamento.
  • kv_cache (List[Dict[str, torch.Tensor]], opcional) — Cache de KV por bloco para inferência autoregressiva. None seleciona o caminho de treinamento.
  • kv_cache_flag (Dict[str, Any], opcional) — Metadados do cache de KV (por exemplo, a flag is_cache_step e as contagens de tokens).
  • attention_mask (BlockMask, opcional) — Máscara causal pré-construída, tipicamente criada via build_attention_mask. Consumida pelos caminhos de treinamento e de pré-preenchimento do KV-cache; o caminho de inferência autoregressiva atende através do KV cache e não utiliza uma máscara completa. Quando None, os caminhos de treinamento / cache constroem a máscara internamente; essa alternativa não é segura para compilação (a função subjacente flex_attention.create_block_mask quebra o grafo sob fullgraph=True), portanto, passe uma máscara pré-construída sempre que envolver forward em torch.compile.
  • attention_kwargs (dict, opcional) — Encaminhado aos processadores de atenção.
  • return_dict (bool, opcional, padrão True) — Se False, retorna tuplas posicionais em vez de uma dataclass de saída.

Retorna

[~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] ou tuple

Quando return_dict é False, uma tuple simples é retornada. Caso contrário, o rollout de treinamento causal (kv_cache is None) retorna um [~models.transformer_2d.Transformer2DModelOutput], enquanto os caminhos de pré-preenchimento de cache e inferência autoregressiva retornam um [AnyFlowFARTransformerOutput].

Passada causal de FAR. Despacha para um dos três caminhos internos:

  • kv_cache is None → rollout de treinamento causal (retorna Transformer2DModelOutput).
  • kv_cache is not None e kv_cache_flag["is_cache_step"] → cache-prefill (retorna AnyFlowFARTransformerOutput com sample=None).
  • Caso contrário → etapa de inferência autoregressiva (retorna AnyFlowFARTransformerOutput).

AnyFlowFARTransformerOutput

classe diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput

< >

( sample: typing.Optional[torch.Tensor] = Nonekv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None )

Parâmetros

  • sample (torch.Tensor ou None) — Alvo de desruido previsto para o chunk autoregressivo. None para o caminho de pré-preenchimento do cache, que apenas grava o cache KV e não produz nenhuma amostra utilizável.
  • kv_cache (list[dict[str, torch.Tensor]], opcional) — Estado do cache KV por bloco usado pelas etapas autoregressivas subsequentes.

Dataclass de saída para os caminhos de encaminhamento causal do AnyFlowFARTransformer3DModel.

Update on GitHub