diffusers/main/vi/api/models/krea2_transformer2d.md

Krea2Transformer2DModel

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

Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

diffusers.Krea2Transformer2DModel[[diffusers.Krea2Transformer2DModel]]

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

Nguồn

Tham số:

in_channels (int, mặc định là 64) : Số kênh tiềm ẩn 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 chiều 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 MLP SwiGLU bên trong mỗi khối.

timestep_embed_dim (int, mặc định là 256) : Độ 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 text encoder mà các hidden states 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 text-encoder đượ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 truy vấn 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 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 trên 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 head-dim trên các trục vị trí rotary (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[[diffusers.Krea2Transformer2DModel.forward]]

forward(hidden_states: Tensor, encoder_hidden_states: Tensor, timestep: Tensor, position_ids: Tensor, encoder_attention_mask: typing.Optional[torch.Tensor] = None, attention_kwargs: dict[str, typing.Any] | None = None, return_dict: bool = True)

Nguồn

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)) : Các trạng thái ẩn của bộ mã hóa văn bản được lấy tại từng lớp, xếp chồng theo 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ó shape (text_seq_len + image_seq_len, 3)) : tọa độ rotary (t, h, w) cho chuỗi kết hợp. Các dòng text đều bằng không; các dòng image chứa tọa độ của lưới latent.

encoder_attention_mask (torch.Tensor có shape (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 truyền tiến.

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.

Trả về:

Transformer2DModelOutput hoặc một tuple mà phần tử đầu tiên là tensor vận tốc có hình dạng (batch_size, image_seq_len, in_channels).

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