Diffusers documentation

Krea2Transformer2DModel

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Krea2Transformer2DModel

Transformer flow-matching MMDiT single-stream được sử dụng bởi Krea 2.

Krea2Transformer2DModel

lớp diffusers.Krea2Transformer2DModel

< >

( in_channels: int = 64num_layers: int = 28attention_head_dim: int = 128num_attention_heads: int = 48num_key_value_heads: int = 12intermediate_size: int = 16384timestep_embed_dim: int = 256text_hidden_dim: int = 2560num_text_layers: int = 12text_num_attention_heads: int = 20text_num_key_value_heads: int = 20text_intermediate_size: int = 6912num_layerwise_text_blocks: int = 2num_refiner_text_blocks: int = 2axes_dims_rope: tuple[int, int, int] = (32, 48, 48)rope_theta: float = 1000.0norm_eps: float = 1e-05 )

Tham số

  • in_channels (int, mặc định là 64) — Số kênh latent sau khi patchification (vae_channels * patch_size ** 2).
  • num_layers (int, mặc định là 28) — Số lượng khối transformer.
  • attention_head_dim (int, mặc định là 128) — Kích thước của mỗi đầu chú ý; kích thước ẩn tổng cộng là attention_head_dim * num_attention_heads.
  • num_attention_heads (int, mặc định là 48) — Số lượng đầu truy vấn.
  • num_key_value_heads (int, mặc định là 12) — Số lượng đầu key/value cho grouped-query attention.
  • intermediate_size (int, mặc định là 16384) — Kích thước ẩn của mạng MLP SwiGLU bên trong mỗi khối.
  • timestep_embed_dim (int, mặc định là 256) — Chiều rộng của nhúng bước thời gian hình sin trước MLP của nó.
  • text_hidden_dim (int, mặc định là 2560) — Kích thước ẩn của bộ mã hóa văn bản mà các trạng thái ẩn của nó được sử dụng.
  • num_text_layers (int, mặc định là 12) — Số lượng trạng thái ẩn của bộ mã hóa văn bản được lấy (tapped) và xếp chồng cho mỗi token.
  • text_num_attention_heads (int, mặc định là 20) — Số lượng đầu query trong các khối kết hợp văn bản.
  • text_num_key_value_heads (int, mặc định là 20) — Số lượng đầu key/value trong các khối kết hợp văn bản.
  • text_intermediate_size (int, mặc định là 6912) — Kích thước ẩn của mạng MLP SwiGLU bên trong các khối kết hợp văn bản.
  • num_layerwise_text_blocks (int, mặc định là 2) — Số lượng khối kết hợp văn bản được áp dụng dọc theo trục lớp được trích (cho mỗi token).
  • num_refiner_text_blocks (int, mặc định là 2) — Số lượng khối kết hợp văn bản được áp dụng trên toàn bộ chuỗi token.
  • axes_dims_rope (tuple[int, int, int], mặc định là (32, 48, 48)) — Cách chia kích thước head trên các trục vị trí xoay (t, h, w).
  • rope_theta (float, mặc định là 1000.0) — Cơ số được sử dụng bởi nhúng vị trí quay (rotary position embedding).
  • norm_eps (float, mặc định là 1e-5) — Epsilon được sử dụng bởi tất cả các module RMSNorm.

Khung xương (backbone) khớp dòng (flow-matching) MMDiT đơn luồng được sử dụng bởi pipeline Krea 2.

Điều kiện văn bản được đưa vào dưới dạng một chồng các trạng thái ẩn được lấy từ nhiều lớp của bộ mã hóa văn bản đa phương thức. Một transformer kết hợp văn bản nhỏ sẽ thu gọn trục lớp và tinh chỉnh chuỗi token; kết quả được nối với các latent ảnh đã được chia patch thành một chuỗi [text, image] duy nhất để xử lý bởi các khối transformer. Timestep điều kiện cho mọi khối thông qua một vector điều chế dùng chung cùng các bảng học được cho từng khối.

forward

< >

( hidden_states: Tensorencoder_hidden_states: Tensortimestep: Tensorposition_ids: Tensorencoder_attention_mask: typing.Optional[torch.Tensor] = Noneattention_kwargs: dict[str, typing.Any] | None = Nonereturn_dict: bool = True )

Tham số

  • hidden_states (torch.Tensor có hình dạng (batch_size, image_seq_len, in_channels)) — Các latent ảnh nhiễu đã được đóng gói (patchified).
  • encoder_hidden_states (torch.Tensor có hình dạng (batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — Ngăn xếp các trạng thái ẩn của bộ mã hóa văn bản được lấy tại từng token.
  • timestep (torch.Tensor có hình dạng (batch_size,)) — Thời gian flow-matching trong [0, 1] (1 là nhiễu thuần túy, 0 là dữ liệu sạch).
  • position_ids (torch.Tensor có hình dạng (text_seq_len + image_seq_len, 3)) — (t, h, w) tọa độ xoay cho chuỗi kết hợp. Các dòng văn bản đều bằng không; các dòng ảnh chứa tọa độ lưới tiềm ẩn.
  • encoder_attention_mask (torch.Tensor có hình dạng (batch_size, text_seq_len), tùy chọn) — Mặt nạ Boolean đánh dấu các token văn bản hợp lệ. Truyền None khi mọi token văn bản đều hợp lệ.
  • attention_kwargs (dict, tùy chọn) — Một từ điển kwargs mà, khi chứa mục scale, sẽ thiết lập tỷ lệ LoRA được áp dụng cho các adapter của transformer này trong suốt quá trình forward pass.
  • return_dict (bool, tùy chọn, mặc định là True) — Có nên trả về một Transformer2DModelOutput thay vì một tuple thông thường hay không.

Dự đoán vận tốc khớp dòng cho các token ảnh.

Update on GitHub