Hugging Face's logo

Diffusers documentation

Flux2Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Flux2Transformer2DModel

Một mô hình Transformer cho dữ liệu dạng ảnh từ Flux2.

Flux2Transformer2DModel

lớp diffusers.Flux2Transformer2DModel

< >

( patch_size: int = 1in_channels: int = 128out_channels: int | None = Nonenum_layers: int = 8num_single_layers: int = 48attention_head_dim: int = 128num_attention_heads: int = 48joint_attention_dim: int = 15360timestep_guidance_channels: int = 256mlp_ratio: float = 3.0axes_dims_rope: tuple[int, ...] = (32, 32, 32, 32)rope_theta: int = 2000eps: float = 1e-06guidance_embeds: bool = True )

Tham số

  • patch_size (int, mặc định là 1) — Kích thước patch để chuyển dữ liệu đầu vào thành các patch nhỏ.
  • in_channels (int, mặc định là 128) — Số lượng kênh trong đầu vào.
  • out_channels (int, tùy chọn, mặc định là None) — Số lượng kênh trong đầu ra. Nếu không được chỉ định, giá trị mặc định sẽ là in_channels.
  • num_layers (int, mặc định là 8) — Số lượng lớp của các khối DiT dual stream cần sử dụng.
  • num_single_layers (int, mặc định là 48) — Số lượng lớp của các khối DiT single stream cần sử dụng.
  • attention_head_dim (int, mặc định là 128) — Số chiều sử dụng cho mỗi đầu chú ý (attention head).
  • num_attention_heads (int, mặc định là 48) — Số lượng đầu attention cần sử dụng.
  • joint_attention_dim (int, mặc định là 15360) — Số chiều được sử dụng cho chú ý liên hợp (chiều nhúng/kênh của encoder_hidden_states).
  • pooled_projection_dim (int, mặc định là 768) — Số chiều sử dụng cho phép chiếu đã gộp (pooled projection).
  • guidance_embeds (bool, mặc định là True) — Có sử dụng guidance embeddings cho biến thể guidance-distilled của mô hình hay không.
  • axes_dims_rope (tuple[int], mặc định là (32, 32, 32, 32)) — Các chiều cần sử dụng cho các nhúng vị trí quay (rotary positional embeddings).

Mô hình Transformer được giới thiệu trong Flux 2.

Tham khảo: https://blackforestlabs.ai/announcing-black-forest-labs/

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensor = Nonetimestep: LongTensor = Noneimg_ids: Tensor = Nonetxt_ids: Tensor = Noneguidance: Tensor = Nonejoint_attention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = Truekv_cache: Flux2KVCache | None = Nonekv_cache_mode: str | None = Nonenum_ref_tokens: int = 0ref_fixed_timestep: float = 0.0 )

Tham số

  • hidden_states (torch.Tensor có hình dạng (batch_size, image_sequence_length, in_channels)) — hidden_states đầu vào.
  • encoder_hidden_states (torch.Tensor có hình dạng (batch_size, text_sequence_length, joint_attention_dim)) — Các embedding điều kiện (các embedding được tính toán từ các điều kiện đầu vào như prompt) cần sử dụng.
  • timestep (torch.LongTensor) — Dùng để chỉ bước khử nhiễu.
  • img_ids (torch.Tensor) — Các id vị trí ảnh được sử dụng để tính các nhúng vị trí quay (rotary positional embeddings).
  • txt_ids (torch.Tensor) — Các id vị trí văn bản được sử dụng để tính các nhúng vị trí quay (rotary positional embeddings).
  • guidance (torch.Tensor, tùy chọn) — Nhúng thang độ dẫn hướng được sử dụng cho các biến thể mô hình đã được chưng cất dẫn hướng.
  • joint_attention_kwargs (dict, tùy chọn) — Một từ điển kwargs, nếu được chỉ định, sẽ được truyền đến AttentionProcessor như đã định nghĩa dưới self.processor trong diffusers.models.attention_processor.
  • return_dict (bool, tùy chọn, mặc định là True) — Có nên trả về một ~models.transformer_2d.Transformer2DModelOutput thay vì một tuple thông thường hay không.
  • kv_cache (Flux2KVCache, tùy chọn) — Bộ nhớ đệm KV cho các token của ảnh tham chiếu. Khi kv_cache_mode là “extract”, một bộ nhớ đệm mới sẽ được tạo và trả về. Khi là “cached”, bộ nhớ đệm được cung cấp sẽ được sử dụng để tiêm K/V của ref trong quá trình attention.
  • kv_cache_mode (str, tùy chọn) — Một trong “extract” (bước đầu tiên với các token tham chiếu) hoặc “cached” (các bước tiếp theo sử dụng K/V tham chiếu đã được lưu cache). Khi None, thực hiện bước truyền tiến chuẩn mà không có lưu cache KV.
  • num_ref_tokens (int, mặc định là 0) — Số lượng token ảnh tham chiếu được thêm vào trước hidden_states (chỉ được sử dụng khi kv_cache_mode="extract").
  • ref_fixed_timestep (float, mặc định là 0.0) — Bước thời gian cố định cho điều biến token tham chiếu (chỉ được sử dụng khi kv_cache_mode="extract").

Phương thức forward của Flux2Transformer2DModel.

Flux2Transformer2DModelOutput

lớp diffusers.models.transformers.transformer_flux2.Flux2Transformer2DModelOutput

< >

( sample: torch.Tensorkv_cache: Flux2KVCache | None = None )

Tham số

  • sample (torch.Tensor có hình dạng (batch_size, num_channels, height, width)) — Trạng thái ẩn đầu ra được điều kiện bởi đầu vào encoder_hidden_states.
  • kv_cache (Flux2KVCache, tùy chọn) — Bộ nhớ đệm KV đã được nạp cho các token của hình ảnh tham chiếu. Chỉ được trả về khi kv_cache_mode="extract".

Đầu ra của Flux2Transformer2DModel.

Update on GitHub

←ErnieImageTransformer2DModel Flux3Transformer2DModel→