Diffusers documentation
AnyFlowFARTransformer3DModel
AnyFlowFARTransformer3DModel
Transformer 3D nhân quả (FAR) được sử dụng bởi AnyFlowFARPipeline —
biến thể FAR của AnyFlow. Xem trang AnyFlowFARPipeline để biết về bài báo, tác giả và các checkpoint đã phát hành. Nó mở rộng
khung xương Wan2.1 phiên bản v0.35.1 với ba bổ sung:
- Mặt nạ khối nhân quả FAR thông qua
torch.nn.attention.flex_attention, hỗ trợ tạo tự hồi quy theo từng khối như đã giới thiệu trong FAR. - Đóng gói patch khung nén (
far_patch_embedding) cho các khung ngữ cảnh (đã được tạo), khởi tạo ấm từpatch_embeddingđộ phân giải đầy đủ tại thời điểm xây dựng thông qua nội suy ba tuyến tính. - Nhúng flow-map hai bước thời gian (giống với
AnyFlowTransformer3DModel) — mỗi lần gọi forward đều điều kiện hóa trên cả bước thời gian nguồntvà bước thời gian đíchr.
Lịch phân chunk mặc định (chunk_partition) được lưu trong cấu hình mô hình; các checkpoint NVIDIA AnyFlow-FAR đã phát hành sử dụng [1, 3, 3, 3, 3, 3, 3, 2] cho cấu hình chuẩn 81 khung hình. forward chấp nhận ghi đè chunk_partition cho từng lần gọi, do đó cùng một checkpoint cũng xử lý được các cấu hình num_frames khác mà không cần huấn luyện lại.
from diffusers import AnyFlowFARTransformer3DModel
# Causal AnyFlow checkpoint (FAR):
transformer = AnyFlowFARTransformer3DModel.from_pretrained(
"nvidia/AnyFlow-FAR-Wan2.1-1.3B-Diffusers", subfolder="transformer"
)AnyFlowFARTransformer3DModel
lớp 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) )
Tham số
- patch_size (Tuple[int], mặc định là (1, 2, 2)) — Kích thước patch 3D cho các khối độ phân giải đầy đủ.
- compressed_patch_size (Tuple[int], mặc định là (1, 4, 4)) — Kích thước patch lớn hơn cho các khối (context) được nén FAR.
- full_chunk_limit (int, mặc định là 3) —
Số lượng tối đa các chunk độ phân giải đầy đủ trước khi các chunk trước đó bị hạ xuống ngữ cảnh FAR nén. Các checkpoint được phát hành sử dụng
3. - num_attention_heads (int, mặc định là 40) — Số lượng đầu chú ý.
- attention_head_dim (int, mặc định là 128) — Số lượng kênh trong mỗi head.
- in_channels (int, mặc định là 16) — Số lượng kênh trong latent đầu vào.
- out_channels (int, mặc định là 16) — Số lượng kênh trong latent đầu ra.
- text_dim (int, mặc định là 4096) — Kích thước đầu vào cho các vector nhúng văn bản (UMT5).
- freq_dim (int, mặc định là 256) — Kích thước cho các nhúng thời gian hình sin.
- ffn_dim (int, mặc định là 13824) — Kích thước chiều trung gian trong mạng feed-forward.
- num_layers (int, mặc định là 40) — Số lượng khối transformer.
- cross_attn_norm (bool, mặc định là True) — Bật chuẩn hóa cross-attention.
- eps (float, mặc định là 1e-6) — Epsilon cho các lớp chuẩn hóa.
- image_dim (Optional[int], optional, mặc định là None) — Kích thước nhúng hình ảnh cho điều kiện I2V.
- rope_max_seq_len (int, mặc định là 1024) — Chiều dài tối đa của chuỗi được sử dụng để tính trước các tần số vị trí quay (rotary position frequencies).
- gate_value (float, mặc định là 0.25) — Cổng trộn giữa các embedding của bước nguồn và bước delta.
- deltatime_type (str, mặc định là ‘r’) —
Hoặc là
"r"(delta là bước thời gian đích) hoặc"t-r"(delta là khoảng thời gian tuyệt đối). - chunk_partition (Tuple[int, …], mặc định là (1, 3, 3, 3, 3, 3, 3, 2)) —
Số lượng khung hình mặc định cho mỗi chunk được pipeline sử dụng. Các checkpoint NVIDIA AnyFlow-FAR đã phát hành nhắm đến
num_frames=81(21 khung hình latent với bước thời gian VAE là 4) được chia thành1 + 3*6 + 2. Một giá trịnum_frameskhác sẽ yêu cầu ghi đèchunk_partitiontương ứng được truyền vào AnyFlowFARPipeline.call() (và tương tự choforward).
Transformer 3D nhân quả (FAR) cho lấy mẫu bản đồ dòng của AnyFlow với tạo tự hồi quy theo từng khối.
Mở rộng bộ khung Wan2.1 phiên bản v0.35.1 với các tính năng sau:
- Mặt nạ khối nhân quả FAR thông qua
torch.nn.attention.flex_attention, hỗ trợ tạo tự hồi quy theo từng khối (FAR). - Đóng gói patch cho khung nén
far_patch_embeddingcho các khung ngữ cảnh (đã được tạo), được khởi tạo từpatch_embeddingthông qua nội suy ba tuyến tính để một mô hình vừa được xây dựng đã có điểm khởi đầu hợp lý ngay cả trước khi tinh chỉnh bằng LoRA. - Nhúng flow-map hai bước thời gian cho lấy mẫu bất kỳ bước nào (giống như
AnyFlowTransformer3DModel).
Hãy dùng AnyFlowTransformer3DModel cho T2V hai chiều thuần túy — biến thể đó bỏ qua FAR causal masking và far_patch_embedding, nhỏ hơn khoảng 5–10%.
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
Tham số
- chunk_partition — số lượng khung hình cho từng chunk (phải có tổng bằng số khung hình latent).
- height, width — các chiều không gian của latent.
- has_clean_context —
Truekhiclean_hidden_statessẽ được truyền quaforward(huấn luyện V2V/I2V); chỉ cờ hiện diện này ảnh hưởng đến bố cục mask. - device — thiết bị cho
BlockMaskkết quả. Mặt nạ không được tự động di chuyển bởidevice_map="auto"; hãy xây dựng nó trên cùng thiết bị mà các đầu vào của transformer sẽ được đặt. - mode —
"train"(khớp với_forward_train) hoặc"cache"(khớp với_forward_cache). Đường dẫn suy luận tự hồi quy_forward_inferencechú ý xuyên qua bộ đệm KV và không có chế độ ở đây.
Trả về
~torch.nn.attention.flex_attention.BlockMask
mặt nạ nhân quả bao phủ bố cục FAR, được đệm đến bội số của 128 dọc theo chiều chuỗi (yêu cầu về kích thước khối của BlockMask).
Gây ra
ValueError
ValueError— nếumodekhông phải là"train"hoặc"cache".
Tạo trước ~torch.nn.attention.flex_attention.BlockMask nhân quả bên ngoài forward.
Truyền kết quả thông qua tham số attention_mask của forward để toàn bộ transformer tương thích với torch.compile(fullgraph=True). Nếu không có mask được xây dựng trước, forward sẽ tự động tạo mask bên trong — con đường này sử dụng flex_attention.create_block_mask(_compile=False) và làm hỏng đồ thị compile.
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
Tham số
- hidden_states (torch.Tensor) —
Đầu vào ẩn có hình dạng
(B, F, C, H, W). - timestep (torch.Tensor) — Bước thời gian t của bản đồ luồng nguồn (nhiễu hơn).
- r_timestep (torch.Tensor) — Bước thời gian r của bản đồ dòng chảy đích (sạch hơn).
- encoder_hidden_states (torch.Tensor) — Nhúng văn bản UMT5.
- chunk_partition (List[int]) —
Số lượng khung hình cho từng chunk; tổng số phải khớp với số khung hình latent trong
hidden_states. - encoder_hidden_states_image (torch.Tensor, optional) — Nhúng ảnh I2V; được nối trước các token văn bản khi được cung cấp.
- clean_hidden_states (torch.Tensor, optional) — Các khung điều kiện sạch (không nhiễu) được sử dụng bởi quá trình rollout trong huấn luyện.
- clean_timestep (torch.Tensor, optional) — Các timestep cho các khung hình điều kiện sạch trong quá trình rollout huấn luyện.
- kv_cache (List[Dict[str, torch.Tensor]], optional) — Bộ nhớ đệm KV theo khối cho suy luận tự hồi quy. None chọn đường dẫn huấn luyện.
- kv_cache_flag (Dict[str, Any], tùy chọn) —
Dữ liệu mô tả của bộ đệm KV (ví dụ cờ
is_cache_stepvà số lượng token). - attention_mask (BlockMask, optional) —
Mặt nạ nhân quả được xây dựng sẵn, thường được tạo thông qua
build_attention_mask. Được sử dụng bởi các đường dẫn huấn luyện và prefill bộ đệm KV; đường dẫn suy luận tự hồi quy sẽ chú ý thông qua bộ đệm KV và không sử dụng mặt nạ đầy đủ. KhiNone, các đường dẫn huấn luyện / bộ đệm sẽ tự động xây dựng mặt nạ bên trong; phương án dự phòng này không an toàn khi biên dịch (hàmflex_attention.create_block_maskbên dưới sẽ phá vỡ đồ thị khi đặtfullgraph=True), do đó hãy truyền một mặt nạ được xây dựng sẵn mỗi khi bọcforwardtrongtorch.compile. - attention_kwargs (dict, optional) — Được chuyển tiếp đến các bộ xử lý attention.
- return_dict (bool, tuỳ chọn, mặc định là True) — Nếu False, trả về các tuple theo vị trí thay vì một dataclass đầu ra.
Trả về
[~models.transformer_2d.Transformer2DModelOutput], [AnyFlowFARTransformerOutput] hoặc tuple
Khi return_dict là False, một tuple thông thường sẽ được trả về. Ngược lại, quy trình rollout huấn luyện causal (kv_cache is None) sẽ trả về [~models.transformer_2d.Transformer2DModelOutput], trong khi các đường dẫn prefill cache và suy luận tự hồi quy sẽ trả về [AnyFlowFARTransformerOutput].
Lượt truyền tiến nhân quả của FAR. Phân phối đến một trong ba đường dẫn nội bộ:
kv_cache is None→ rollout huấn luyện causal (trả vềTransformer2DModelOutput).kv_cache is not Nonevàkv_cache_flag["is_cache_step"]→ cache-prefill (trả vềAnyFlowFARTransformerOutputvớisample=None).- Ngược lại → bước suy luận tự hồi quy (trả về
AnyFlowFARTransformerOutput).
AnyFlowFARTransformerOutput
lớp 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 )
Tham số
- sample (torch.Tensor hoặc None) —
Mục tiêu khử nhiễu dự đoán cho khối tự hồi quy.
Nonecho đường dẫn prefill cache, chỉ ghi vào KV cache và không tạo ra mẫu nào có thể sử dụng được. - kv_cache (list[dict[str, torch.Tensor]], optional) — Trạng thái bộ đệm KV theo khối được sử dụng bởi các bước tự hồi quy tiếp theo.
Dataclass đầu ra cho các đường dẫn tiến hành nhân quả của AnyFlowFARTransformer3DModel.