Diffusers documentation
Krea2Transformer2DModel
Krea2Transformer2DModel
Krea 2에서 사용하는 단일 스트림 MMDiT 플로우 매칭 트랜스포머입니다.
Krea2Transformer2DModel
클래스 diffusers.Krea2Transformer2DModel
< source >( 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 = (32, 48, 48)rope_theta: float = 1000.0norm_eps: float = 1e-05 )
매개변수
- in_channels (
int, 기본값 64) — 패치화 이후의 Latent 채널 수(vae_channels * patch_size ** 2). - num_layers (
int, 기본값 28) — Transformer 블록의 수. - attention_head_dim (
int, 기본값 128) — 각 attention head의 차원입니다. 전체 hidden size는attention_head_dim * num_attention_heads입니다. - num_attention_heads (
int, 기본값 48) — 쿼리 헤드 수. - num_key_value_heads (
int, 기본값 12) — grouped-query attention에 사용할 key/value head의 수입니다. - intermediate_size (
int, 기본값 16384) — 각 블록 내부의 SwiGLU MLP의 피드포워드 히든 크기입니다. - timestep_embed_dim (
int, 기본값 256) — MLP 적용 전 사인oidal timestep 임베딩의 폭. - text_hidden_dim (
int, 기본값 2560) — 소비되는 hidden states를 가진 텍스트 인코더의 hidden size입니다. - num_text_layers (
int, 기본값 12) — 토큰당 스택되는 텍스트 인코더의 탭된 히든 상태 수. - text_num_attention_heads (
int, 기본값 20) — 텍스트 융합 블록의 쿼리 헤드 수입니다. - text_num_key_value_heads (
int, 기본값 20) — 텍스트 퓨전 블록의 key/value 헤드 수입니다. - text_intermediate_size (
int, 기본값 6912) — 텍스트 융합 블록 내부의 SwiGLU MLP 피드포워드 히든 크기입니다. - num_layerwise_text_blocks (
int, 기본값 2) — 토큰별로 적용되는 tapped-layer 축에 걸쳐 사용할 텍스트 퓨전 블록의 수입니다. - num_refiner_text_blocks (
int, 기본값 2) — 토큰 시퀀스에 걸쳐 적용되는 텍스트 융합 블록의 수입니다. - axes_dims_rope (
tuple[int, int, int], 기본값(32, 48, 48)) — (t, h, w) 회전 위치 축에 걸쳐 head-dim을 분할하는 값입니다. - rope_theta (
float, 기본값 1000.0) — 회전 위치 임베딩에 사용되는 기본값입니다. - norm_eps (
float, 기본값 1e-5) — 모든 RMSNorm 모듈에서 사용되는 Epsilon 값입니다.
Krea 2 파이프라인에서 사용하는 단일 스트림 MMDiT 플로우 매칭 백본입니다.
Text conditioning은 멀티모달 텍스트 인코더의 여러 레이어에서 추출된 hidden states 스택으로 입력됩니다. 작은 text-fusion transformer가 레이어 축을 축소하고 토큰 시퀀스를 정제하며, 그 결과는 패치화된 이미지 latents와 결합되어 단일 [text, image] 시퀀스를 형성하고 transformer 블록에 의해 처리됩니다. timestep은 하나의 공유 modulation 벡터와 블록별 학습된 테이블을 통해 모든 블록을 조건화합니다.
forward
< source >( 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 )
매개변수
- hidden_states (
torch.Tensor, shape(batch_size, image_seq_len, in_channels)) — 패키징(패치화)된 노이즈 이미지 레이턴트입니다. - encoder_hidden_states (
torch.Tensor, shape(batch_size, text_seq_len, num_text_layers, text_hidden_dim)) — 토큰별 텍스트 인코더의 탭된(hidden states) 스택. - timestep (
torch.Tensor, shape(batch_size,)) —[0, 1]범위의 Flow-matching 시간 (1은 순수 노이즈, 0은 깨끗한 데이터를 의미합니다). - position_ids (
torch.Tensor, shape(text_seq_len + image_seq_len, 3)) — 결합된 시퀀스에 대한(t, h, w)rotary 좌표입니다. 텍스트 행은 모두 0이며, 이미지 행에는 latent-grid 좌표가 저장됩니다. - encoder_attention_mask (
torch.Tensor, shape(batch_size, text_seq_len), 선택 사항) — 유효한 텍스트 토큰을 표시하는 불리언 마스크입니다. 모든 텍스트 토큰이 유효한 경우None을 전달하세요. - attention_kwargs (
dict, 선택 사항) —scale항목을 포함할 경우, forward pass 동안 이 transformer의 어댑터에 적용되는 LoRA scale을 설정하는 kwargs 사전입니다. - return_dict (
bool, 선택 사항, 기본값True) — 일반 튜플 대신 Transformer2DModelOutput을 반환할지 여부.
이미지 토큰에 대한 flow-matching 속도를 예측합니다.