Diffusers documentation

AnyFlowFARTransformer3DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

AnyFlowFARTransformer3DModel

AnyFlowFARPipeline에서 사용하는 인과적(causal) FAR 3D Transformer — AnyFlow의 FAR 변형입니다. 논문, 저자, 공개된 체크포인트에 대한 자세한 내용은 AnyFlowFARPipeline 페이지를 참고하세요. v0.35.1 Wan2.1 백본을 기반으로 다음 세 가지를 추가하여 확장했습니다:

  1. FAR에서 소개한 청크 단위 자기회귀 생성을 지원하는 torch.nn.attention.flex_attention를 통한 FAR causal block-mask
  2. 컨텍스트(이미 생성된) 프레임에 대한 압축 프레임 패치 임베딩(far_patch_embedding)은 생성 시점에 삼선형 보간을 통해 풀 해상도 patch_embedding에서 웜 스타트됩니다.
  3. 이중 타임스텝 플로우맵 임베딩(AnyFlowTransformer3DModel과 동일) — 모든 forward 호출은 소스 타임스텝 t와 대상 타임스텝 r 모두에 조건을 부여합니다.

기본 chunk schedule(chunk_partition)은 모델 config에 저장됩니다. 공개된 NVIDIA AnyFlow-FAR 체크포인트는 표준 81프레임 설정에 대해 [1, 3, 3, 3, 3, 3, 3, 2]를 사용합니다. forward는 호출별 chunk_partition 오버라이드를 지원하므로, 동일한 체크포인트로 재학습 없이 다른 num_frames 구성도 처리할 수 있습니다.

from diffusers import AnyFlowFARTransformer3DModel

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

AnyFlowFARTransformer3DModel

클래스 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) )

매개변수

  • patch_size (Tuple[int], 기본값 (1, 2, 2)) — 풀 해상도 청크에 대한 3D 패치 차원입니다.
  • compressed_patch_size (Tuple[int], 기본값 (1, 4, 4)) — FAR 압축(컨텍스트) 청크에 사용할 더 큰 패치 차원입니다.
  • full_chunk_limit (int, 기본값 3) — 이전 chunk가 압축된 FAR 컨텍스트로 강등되기 전까지 허용되는 최대의 full-resolution chunk 수입니다. 공개된 체크포인트는 3을 사용합니다.
  • num_attention_heads (int, 기본값 40) — Attention head의 수입니다.
  • attention_head_dim (int, 기본값 128) — 각 헤드당 채널 수입니다.
  • in_channels (int, 기본값 16) — 입력 latent의 채널 수입니다.
  • out_channels (int, 기본값 16) — 출력 latent의 채널 수입니다.
  • text_dim (int, 기본값 4096) — 텍스트 임베딩(UMT5)의 입력 차원입니다.
  • freq_dim (int, 기본값 256) — 사인 파 시간 임베딩에 사용되는 차원입니다.
  • ffn_dim (int, 기본값 13824) — 피드포워드 네트워크의 중간 차원입니다.
  • num_layers (int, 기본값 40) — Transformer 블록의 수입니다.
  • cross_attn_norm (bool, 기본값 True) — 크로스 어텐션 정규화를 활성화합니다.
  • eps (float, 기본값 1e-6) — 정규화 레이어에 사용되는 Epsilon 값입니다.
  • image_dim (Optional[int], 선택 사항, 기본값 None) — I2V 조건화에 사용되는 이미지 임베딩 차원입니다.
  • rope_max_seq_len (int, 기본값 1024) — 회전 위치 주파수를 사전 계산하는 데 사용되는 최대 시퀀스 길이입니다.
  • gate_value (float, 기본값 0.25) — 소스 타임스텝 임베딩과 델타 타임스텝 임베딩 사이의 혼합 게이트입니다.
  • deltatime_type (str, 기본값 ‘r’) — "r"(delta가 대상 timestep임) 또는 "t-r"(delta가 절대 간격임).
  • chunk_partition (Tuple[int, …], 기본값 (1, 3, 3, 3, 3, 3, 3, 2)) — 파이프라인에서 사용하는 기본 청크별 프레임 수입니다. 공개된 NVIDIA AnyFlow-FAR 체크포인트는 num_frames=81 (VAE 시간 스트라이드 4 기준 21개의 잠재 프레임)을 1 + 3*6 + 2로 분할하여 대상으로 합니다. 다른 num_frames를 사용하려면 chunk_partition 오버라이드를 AnyFlowFARPipeline.call() (그리고 forward에도 동일하게)에 전달해야 합니다.

AnyFlow 플로우 맵 샘플링을 위한 청크 단위 오토레거시 생성을 지원하는 인과적(FAR) 3D Transformer.

v0.35.1 Wan2.1 백본을 확장하여 다음을 추가합니다:

  • torch.nn.attention.flex_attention를 통한 FAR causal block-mask로, 청크 단위 자기회귀 생성(FAR)을 지원합니다.
  • 컨텍스트(이미 생성된) 프레임에 대한 압축 프레임 패치 임베딩 far_patch_embedding은 patch_embedding에서 삼선형 보간을 통해 초기화되어, LoRA 미세 조정 전에 새로 생성된 모델도 이미 합리적인 시작점에 있게 됩니다.
  • AnyFlowTransformer3DModel과 동일한 any-step 샘플링을 위한 Dual-timestep flow-map embedding

순수 양방향 T2V에는 AnyFlowTransformer3DModel을 대신 사용하세요. 해당 변형은 FAR 인과적 마스킹과 far_patch_embedding을 생략하여 크기가 약 5–10% 더 작습니다.

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

매개변수

  • chunk_partition — 청크별 프레임 수 (잠재 프레임 수의 합과 같아야 함).
  • height, width — latent 공간 차원.
  • has_clean_context — clean_hidden_states가 forward를 통해 전달되는 경우(V2V/I2V 학습) True입니다. 마스크 레이아웃에 영향을 미치는 것은 이 존재 플래그뿐입니다.
  • device — 결과 BlockMask에 사용할 디바이스입니다. device_map="auto"로 마스크가 자동으로 이동되지 않으므로, 트랜스포머의 입력이 존재할 디바이스와 동일한 디바이스에서 생성해야 합니다.
  • mode — "train"(_forward_train와 일치) 또는 "cache"(_forward_cache와 일치). 자기 회귀 _forward_inference 경로는 KV 캐시를 통해 주의를 기울이며 여기에는 mode가 없습니다.

반환값

~torch.nn.attention.flex_attention.BlockMask

FAR 레이아웃을 가로지르는 causal mask로, 시퀀스 차원(시퀀스 길이)을 128의 배수로 패딩합니다(BlockMask의 블록 크기 요구 사항).

발생

ValueError

  • ValueError — mode가 "train"도 "cache"도 아닌 경우.

forward 외부에서 인과적 ~torch.nn.attention.flex_attention.BlockMask를 사전 생성하세요.

forward의 attention_mask kwarg를 통해 결과를 전달하여 전체 transformer를 torch.compile(fullgraph=True)와 호환되도록 합니다. 미리 생성된 마스크가 없으면 forward는 내부적으로 마스크를 구성하는 방식으로 폴백됩니다. 이 경로는 flex_attention.create_block_mask(_compile=False)를 사용하며 컴파일 그래프를 깨뜨립니다.

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

매개변수

  • hidden_states (torch.Tensor) — (B, F, C, H, W) 형태의 잠재 입력(latent input)
  • timestep (torch.Tensor) — 소스(노이즈가 더 많은) flow-map timestep t.
  • r_timestep (torch.Tensor) — 목표(더 깨끗한) flow-map timestep r.
  • encoder_hidden_states (torch.Tensor) — UMT5 텍스트 임베딩.
  • chunk_partition (List[int]) — 각 청크의 프레임 수. 총합은 hidden_states의 잠재 프레임 수와 일치해야 합니다.
  • encoder_hidden_states_image (torch.Tensor, optional) — I2V 이미지 임베딩입니다. 제공된 경우 텍스트 토큰 앞에 연결됩니다.
  • clean_hidden_states (torch.Tensor, optional) — 학습 롤아웃에 사용되는 노이즈가 없는(conditioning) 프레임입니다.
  • clean_timestep (torch.Tensor, optional) — 학습 rollout에서 clean conditioning 프레임에 사용되는 timesteps입니다.
  • kv_cache (List[Dict[str, torch.Tensor]], 선택 사항) — 자기 회귀 추론을 위한 블록별 KV 캐시입니다. None을 선택하면 학습 경로를 사용합니다.
  • kv_cache_flag (Dict[str, Any], 선택 사항) — KV-cache 메타데이터(예: is_cache_step 플래그 및 토큰 수).
  • attention_mask (BlockMask, 선택) — build_attention_mask를 통해 일반적으로 생성되는 사전 구축된 causal mask입니다. 학습 및 KV-cache prefill 경로에서 사용되며, 자기 회귀 추론 경로는 KV cache를 통해 attention을 수행하므로 전체 mask를 사용하지 않습니다. None인 경우 학습/캐시 경로가 내부적으로 mask를 생성하지만, 해당 폴백은 컴파일 안전하지 않습니다(기저의 flex_attention.create_block_mask가 fullgraph=True 환경에서 그래프를 깨뜨림). 따라서 forward를 torch.compile로 감싸는 경우 항상 사전 구축된 mask를 전달해야 합니다.
  • attention_kwargs (dict, 선택 사항) — attention processor로 전달됩니다.
  • return_dict (bool, 선택 사항, 기본값 True) — False인 경우 출력 dataclass 대신 위치 기반 튜플을 반환합니다.

반환값

[~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] 또는 tuple

return_dict가 False인 경우, 일반 tuple이 반환됩니다. 그렇지 않으면, causal training rollout(kv_cache is None)은 [~models.transformer_2d.Transformer2DModelOutput]을 반환하고, cache-prefill 및 autoregressive inference 경로는 [AnyFlowFARTransformerOutput]을 반환합니다.

FAR 인과적(forward) 패스. 세 가지 내부 경로 중 하나로 분기합니다:

  • kv_cache is None → causal training rollout(Transformer2DModelOutput 반환).
  • kv_cache is not None이고 kv_cache_flag["is_cache_step"]인 경우 → cache-prefill(sample=None인 AnyFlowFARTransformerOutput 반환).
  • 그 외 → 자기회귀 추론 단계(AnyFlowFARTransformerOutput 반환).

AnyFlowFARTransformerOutput

class diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput

< >

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

매개변수

  • sample (torch.Tensor 또는 None) — 자기 회귀 청크에 대한 예측된 디노이징 대상입니다. 캐시 프리필 경로의 경우 None이며, 이 경로는 KV 캐시만 작성하고 사용 가능한 샘플을 생성하지 않습니다.
  • kv_cache (list[dict[str, torch.Tensor]], optional) — 이후의 자기회귀 단계에서 사용되는 블록별 KV 캐시 상태.

AnyFlowFARTransformer3DModel의 인과적(forward) 경로에 대한 출력 dataclass입니다.

Update on GitHub