Hugging Face's logo

Diffusers documentation

AnyFlowFARTransformer3DModel

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

←AllegroTransformer3DModel AnyFlowTransformer3DModel→