Buckets:

hf-doc-build/doc / diffusers /main /ko /api /models /anyflow_far_transformer3d.md
mishig's picture
|
download
raw
12.9 kB

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 checkpoints는 표준 81-frame 설정에 대해 [1, 3, 3, 3, 3, 3, 3, 2]를 사용합니다. forward는 호출별 chunk_partition 오버라이드를 지원하므로, 동일한 checkpoint는 재학습 없이 다른 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]]

diffusers.AnyFlowFARTransformer3DModel[[diffusers.AnyFlowFARTransformer3DModel]]

diffusers.AnyFlowFARTransformer3DModel(patch_size: typing.Tuple[int] = (1, 2, 2), compressed_patch_size: typing.Tuple[int] = (1, 4, 4), full_chunk_limit: int = 3, num_attention_heads: int = 40, attention_head_dim: int = 128, in_channels: int = 16, out_channels: int = 16, text_dim: int = 4096, freq_dim: int = 256, ffn_dim: int = 13824, num_layers: int = 40, cross_attn_norm: bool = True, eps: float = 1e-06, image_dim: typing.Optional[int] = None, rope_max_seq_len: int = 1024, gate_value: float = 0.25, deltatime_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) : 이전 청크가 압축된 FAR 컨텍스트로 강등되기 전까지 허용되는 최대 풀 해상도 청크 수. 공개된 체크포인트는 3을 사용합니다.

num_attention_heads (int, 기본값 40) : 어텐션 헤드 수.

attention_head_dim (int, 기본값 128) : 각 헤드(head)의 채널 수.

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) : cross-attention 정규화를 활성화합니다.

eps (float, 기본값 1e-6) : 정규화 레이어에 사용되는 Epsilon 값입니다.

image_dim (Optional[int], 선택 사항, 기본값 None) : I2V 조건화에 사용되는 이미지 임베딩 차원입니다.

rope_max_seq_len (int, 기본값 1024) : 회전 위치 주파수를 사전 계산하는 데 사용되는 최대 시퀀스 길이입니다.

gate_value (float, 기본값 0.25) : source-timestep 임베딩과 delta-timestep 임베딩 사이의 혼합 게이트.

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를 사용하려면 AnyFlowFARPipeline.call() (그리고 forward에도 동일하게)에 전달되는 chunk_partition 오버라이드가 필요합니다.

Causal (FAR) AnyFlow 플로우 맵 샘플링을 위한 청크 단위 자기 회귀 생성을 지원하는 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[[diffusers.AnyFlowFARTransformer3DModel.build_attention_mask]]

build_attention_mask(chunk_partition: typing.List[int], height: int, width: int, has_clean_context: bool = False, device: typing.Optional[torch.device] = None, mode: str = 'train')

소스

매개변수:

chunk_partition : 청크별 프레임 수 (잠재 프레임 수의 합과 같아야 함).

height, width : latent 공간 차원.

has_clean_context : clean_hidden_states가 forward를 통해 전달되는 경우(V2V/I2V 학습) True입니다. 마스크 레이아웃에는 이 존재 플래그만 영향을 줍니다.

device : 결과 BlockMask의 device입니다. device_map="auto"에 의해 마스크가 자동으로 이동되지 않으므로, transformer의 입력이 존재할 device와 동일한 device에서 생성해야 합니다.

mode : "train" (_forward_train와 일치) 또는 "cache" (_forward_cache와 일치). 자기 회귀 _forward_inference 경로는 KV 캐시를 통해 주의를 기울이며 여기에는 모드가 없습니다.

반환값: ~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[[diffusers.AnyFlowFARTransformer3DModel.forward]]

forward(hidden_states: Tensor, timestep: Tensor, r_timestep: Tensor, encoder_hidden_states: Tensor, chunk_partition: typing.List[int], encoder_hidden_states_image: typing.Optional[torch.Tensor] = None, clean_hidden_states: typing.Optional[torch.Tensor] = None, clean_timestep: typing.Optional[torch.Tensor] = None, kv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None, kv_cache_flag: typing.Optional[typing.Dict[str, typing.Any]] = None, attention_mask: typing.Optional[torch.nn.attention.flex_attention.BlockMask] = None, attention_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = None, return_dict: bool = True)

소스

매개변수:

hidden_states (torch.Tensor) : (B, F, C, H, W) 형태의 Latent 입력.

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, 선택) : I2V 이미지 임베딩; 제공되면 텍스트 토큰 앞에 연결됩니다.

clean_hidden_states (torch.Tensor, 선택 사항) : 학습 롤아웃에서 사용되는 노이즈가 없는 클린 조건 프레임입니다.

clean_timestep (torch.Tensor, 선택) : 학습 롤아웃에서 깨끗한 조건 프레임의 Timesteps입니다.

kv_cache (List[Dict[str, torch.Tensor]], optional) : 자기 회귀 추론을 위한 블록별 KV 캐시입니다. None을 선택하면 학습 경로를 사용합니다.

kv_cache_flag (Dict[str, Any], 선택) : KV-cache 메타데이터(예: is_cache_step 플래그 및 토큰 수).

attention_mask (BlockMask, optional) : build_attention_mask를 통해 일반적으로 생성되는 사전 구축된 causal mask입니다. 학습 및 KV-cache prefill 경로에서 사용되며, 자기 회귀 추론 경로는 KV cache를 통해 attention을 수행하므로 전체 mask를 사용하지 않습니다. None인 경우 학습 / cache 경로가 내부적으로 mask를 생성하지만, 해당 폴백은 compile-safe하지 않습니다(기저의 flex_attention.create_block_mask가 fullgraph=True 환경에서 그래프를 깨뜨림). 따라서 forward를 torch.compile로 감싸는 경우 항상 사전 구축된 mask를 전달해야 합니다.

attention_kwargs (dict, optional) : attention processors로 전달됩니다.

return_dict (bool, optional, 기본값 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[[diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput]]

diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput[[diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput]]

diffusers.models.transformers.transformer_anyflow_far.AnyFlowFARTransformerOutput(sample: typing.Optional[torch.Tensor] = None, kv_cache: typing.Optional[typing.List[typing.Dict[str, torch.Tensor]]] = None)

소스

매개변수:

sample (torch.Tensor 또는 None) : 자기 회귀 청크에 대한 예측된 디노이징 대상. KV 캐시만 작성하고 사용 가능한 샘플을 생성하지 않는 캐시 프리필 경로에서는 None입니다.

kv_cache (list[dict[str, torch.Tensor]], 선택) : 이후의 자기회귀 단계에서 사용되는 블록별 KV 캐시 상태입니다.

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

Xet Storage Details

Size:
12.9 kB
·
Xet hash:
0be7cbe1a896666e2d9207596862e4b48f048f46e7c7190150d08dcfba2a4de1

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.