Diffusers documentation
AnyFlowFARTransformer3DModel
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:
- máscara de bloco causal FAR via
torch.nn.attention.flex_attention, com suporte à geração autoregressiva por blocos, conforme introduzido em FAR. - Embedding de patches de frames comprimidos (
far_patch_embedding) para frames de contexto (já gerados), inicializado a partir dopatch_embeddingde resolução completa no momento da construção, por interpolação trilinear. - Embedding de fluxo-mapa com dois passos de tempo (igual ao
AnyFlowTransformer3DModel) — cada chamada de forward é condicionada tanto no passo de tempo de origemtquanto no passo de tempo de destinor.
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
< source >( 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 como1 + 3*6 + 2. Umnum_framesdiferente exige uma substituição correspondente dechunk_partitionpassada para AnyFlowFARPipeline.call() (e da mesma forma paraforward).
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_embeddingpara frames de contexto (já gerados), inicializado a partir depatch_embeddingvia 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
< source >( 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 —
Truequandoclean_hidden_statesserá passado porforward(treino de V2V/I2V); apenas essa flag de presença afeta o layout da máscara. - device — dispositivo para o
BlockMaskresultante. A máscara não é movida automaticamente pordevice_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_inferenceatende 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— semodenã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
< source >( 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_stepe 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. QuandoNone, os caminhos de treinamento / cache constroem a máscara internamente; essa alternativa não é segura para compilação (a função subjacenteflex_attention.create_block_maskquebra o grafo sobfullgraph=True), portanto, passe uma máscara pré-construída sempre que envolverforwardemtorch.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 (retornaTransformer2DModelOutput).kv_cache is not Noneekv_cache_flag["is_cache_step"]→ cache-prefill (retornaAnyFlowFARTransformerOutputcomsample=None).- Caso contrário → etapa de inferência autoregressiva (retorna
AnyFlowFARTransformerOutput).
AnyFlowFARTransformerOutput
classe diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput
< source >( 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.
Nonepara 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.