Skip to content

vllm_omni.diffusion.models.krea2.krea2_transformer

logger module-attribute 

logger = init_logger(__name__)

Krea2Attention

Bases: Module

Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate.

Q/K/V layout is [B, seq, heads, head_dim]; attention_mask is a 2D boolean key-padding mask (batch, key_seq_len) that the backend broadcasts.

attn instance-attribute 

attn = Attention(
    num_heads=self.num_heads,
    head_size=self.head_dim,
    softmax_scale=1.0 / self.head_dim**0.5,
    causal=False,
    num_kv_heads=self.num_kv_heads,
    role=role,
)

head_dim instance-attribute 

head_dim = hidden_size // num_heads

hidden_size instance-attribute 

hidden_size = hidden_size

norm_k instance-attribute 

norm_k = Krea2RMSNorm(self.head_dim, eps=eps)

norm_q instance-attribute 

norm_q = Krea2RMSNorm(self.head_dim, eps=eps)

num_heads instance-attribute 

num_heads = num_heads

num_kv_heads instance-attribute 

num_kv_heads = (
    num_kv_heads if num_kv_heads is not None else num_heads
)

to_gate instance-attribute 

to_gate = _linear(
    hidden_size,
    hidden_size,
    False,
    quant_config,
    _join_prefix(prefix, "to_gate"),
)

to_k instance-attribute 

to_k = _linear(
    hidden_size,
    kv_dim,
    False,
    quant_config,
    _join_prefix(prefix, "to_k"),
)

to_out instance-attribute 

to_out = nn.ModuleList(
    [
        _linear(
            hidden_size,
            hidden_size,
            False,
            quant_config,
            _join_prefix(prefix, "to_out.0"),
        ),
        nn.Dropout(0.0),
    ]
)

to_q instance-attribute 

to_q = _linear(
    hidden_size,
    q_dim,
    False,
    quant_config,
    _join_prefix(prefix, "to_q"),
)

to_v instance-attribute 

to_v = _linear(
    hidden_size,
    kv_dim,
    False,
    quant_config,
    _join_prefix(prefix, "to_v"),
)

forward 

forward(
    hidden_states: Tensor,
    attention_mask: Tensor | None = None,
    image_rotary_emb: tuple[Tensor, Tensor] | None = None,
) -> Tensor

Krea2FinalLayer

Bases: Module

Final adaptive RMSNorm and output projection.

linear instance-attribute 

linear = _linear(
    hidden_size,
    out_channels,
    True,
    quant_config,
    _join_prefix(prefix, "linear"),
)

norm instance-attribute 

norm = Krea2RMSNorm(hidden_size, eps=eps)

scale_shift_table instance-attribute 

scale_shift_table = nn.Parameter(
    torch.zeros(2, hidden_size)
)

forward 

forward(hidden_states: Tensor, temb: Tensor) -> Tensor

Krea2RMSNorm

Bases: Module

RMSNorm with a zero-centered scale: the effective multiplier is 1 + weight, matching the Krea 2 checkpoint format. Activations are upcast so the normalization runs in float32.

dim instance-attribute 

dim = dim

eps instance-attribute 

eps = eps

weight instance-attribute 

weight = nn.Parameter(torch.zeros(dim))

forward 

forward(hidden_states: Tensor) -> Tensor

Krea2RotaryPosEmbed

Bases: Module

Multi-axis (t, h, w) rotary position embedding, following the Flux/Krea 2 convention.

axes_dim instance-attribute 

axes_dim = axes_dim

theta instance-attribute 

theta = theta

forward 

forward(ids: Tensor) -> tuple[Tensor, Tensor]

Krea2SwiGLU

Bases: Module

SwiGLU feed-forward network.

down instance-attribute 

down = _linear(
    hidden_dim,
    dim,
    False,
    quant_config,
    _join_prefix(prefix, "down"),
)

gate instance-attribute 

gate = _linear(
    dim,
    hidden_dim,
    False,
    quant_config,
    _join_prefix(prefix, "gate"),
)

up instance-attribute 

up = _linear(
    dim,
    hidden_dim,
    False,
    quant_config,
    _join_prefix(prefix, "up"),
)

forward 

forward(hidden_states: Tensor) -> Tensor

Krea2TextFusion

Bases: Module

Fuses the stack of tapped text-encoder hidden states into a single sequence of text features.

Two layerwise_blocks attend across the num_text_layers axis independently for every token, a linear projector collapses that axis, and two refiner_blocks attend across the token sequence.

layerwise_blocks instance-attribute 

layerwise_blocks = nn.ModuleList(
    [
        (
            Krea2TextFusionBlock(
                dim,
                num_heads,
                num_kv_heads,
                intermediate_size,
                eps,
                quant_config=quant_config,
                prefix=_join_prefix(
                    prefix, f"layerwise_blocks.{i}"
                ),
            )
        )
        for i in (range(num_layerwise_blocks))
    ]
)

projector instance-attribute 

projector = _linear(
    num_text_layers,
    1,
    False,
    quant_config,
    _join_prefix(prefix, "projector"),
)

refiner_blocks instance-attribute 

refiner_blocks = nn.ModuleList(
    [
        (
            Krea2TextFusionBlock(
                dim,
                num_heads,
                num_kv_heads,
                intermediate_size,
                eps,
                quant_config=quant_config,
                prefix=_join_prefix(
                    prefix, f"refiner_blocks.{i}"
                ),
            )
        )
        for i in (range(num_refiner_blocks))
    ]
)

forward 

forward(
    encoder_hidden_states: Tensor,
    attention_mask: Tensor | None = None,
) -> Tensor

Krea2TextFusionBlock

Bases: Module

Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage.

attn instance-attribute 

attn = Krea2Attention(
    dim,
    num_heads,
    num_kv_heads,
    eps=eps,
    role="self",
    quant_config=quant_config,
    prefix=_join_prefix(prefix, "attn"),
)

ff instance-attribute 

ff = Krea2SwiGLU(
    dim,
    intermediate_size,
    quant_config=quant_config,
    prefix=_join_prefix(prefix, "ff"),
)

norm1 instance-attribute 

norm1 = Krea2RMSNorm(dim, eps=eps)

norm2 instance-attribute 

norm2 = Krea2RMSNorm(dim, eps=eps)

forward 

forward(
    hidden_states: Tensor,
    attention_mask: Tensor | None = None,
) -> Tensor

Krea2TextProjection

Bases: Module

Projects the fused text features into the transformer width.

linear_1 instance-attribute 

linear_1 = _linear(
    text_dim,
    hidden_size,
    True,
    quant_config,
    _join_prefix(prefix, "linear_1"),
)

linear_2 instance-attribute 

linear_2 = _linear(
    hidden_size,
    hidden_size,
    True,
    quant_config,
    _join_prefix(prefix, "linear_2"),
)

norm instance-attribute 

norm = Krea2RMSNorm(text_dim, eps=eps)

forward 

forward(hidden_states: Tensor) -> Tensor

Krea2TimestepEmbedding

Bases: Module

Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP.

Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens.

embed_dim instance-attribute 

embed_dim = embed_dim

linear_1 instance-attribute 

linear_1 = _linear(
    embed_dim,
    hidden_size,
    True,
    quant_config,
    _join_prefix(prefix, "linear_1"),
)

linear_2 instance-attribute 

linear_2 = _linear(
    hidden_size,
    hidden_size,
    True,
    quant_config,
    _join_prefix(prefix, "linear_2"),
)

forward 

forward(timestep: Tensor, dtype: dtype) -> Tensor

Krea2Transformer2DModel

Bases: Module

The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline (vLLM-Omni port).

Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated with the patchified image latents into a single [text, image] sequence processed by the transformer blocks. The timestep conditions every block through one shared modulation vector plus per-block learned tables.

dtype instance-attribute 

dtype = (
    od_config.dtype
    if od_config is not None
    else torch.get_default_dtype()
)

final_layer instance-attribute 

final_layer = Krea2FinalLayer(
    hidden_size,
    in_channels,
    norm_eps,
    quant_config,
    "final_layer",
)

hidden_size instance-attribute 

hidden_size = hidden_size

img_in instance-attribute 

img_in = _linear(
    in_channels, hidden_size, True, quant_config, "img_in"
)

in_channels instance-attribute 

in_channels = in_channels

num_layers instance-attribute 

num_layers = num_layers

od_config instance-attribute 

od_config = od_config

out_channels instance-attribute 

out_channels = in_channels

rotary_emb instance-attribute 

rotary_emb = Krea2RotaryPosEmbed(
    theta=rope_theta, axes_dim=list(axes_dims_rope)
)

text_fusion instance-attribute 

text_fusion = Krea2TextFusion(
    num_text_layers=num_text_layers,
    dim=text_hidden_dim,
    num_heads=text_num_attention_heads,
    num_kv_heads=text_num_key_value_heads,
    intermediate_size=text_intermediate_size,
    num_layerwise_blocks=num_layerwise_text_blocks,
    num_refiner_blocks=num_refiner_text_blocks,
    eps=norm_eps,
    quant_config=quant_config,
    prefix="text_fusion",
)

time_embed instance-attribute 

time_embed = Krea2TimestepEmbedding(
    timestep_embed_dim,
    hidden_size,
    quant_config,
    "time_embed",
)

time_mod_proj instance-attribute 

time_mod_proj = _linear(
    hidden_size,
    6 * hidden_size,
    True,
    quant_config,
    "time_mod_proj",
)

transformer_blocks instance-attribute 

transformer_blocks = nn.ModuleList(
    [
        (
            Krea2TransformerBlock(
                hidden_size=hidden_size,
                intermediate_size=intermediate_size,
                num_heads=num_attention_heads,
                num_kv_heads=num_key_value_heads,
                norm_eps=norm_eps,
                quant_config=quant_config,
                prefix=f"transformer_blocks.{i}",
            )
        )
        for i in (range(num_layers))
    ]
)

txt_in instance-attribute 

txt_in = Krea2TextProjection(
    text_hidden_dim,
    hidden_size,
    norm_eps,
    quant_config,
    "txt_in",
)

forward 

forward(
    hidden_states: Tensor,
    encoder_hidden_states: Tensor,
    timestep: Tensor,
    position_ids: Tensor,
    encoder_attention_mask: Tensor | None = None,
) -> Tensor

Predict the flow-matching velocity for the image tokens.

Parameters:

Name Type Description Default
hidden_states Tensor

Packed (patchified) noisy image latents (batch, image_seq_len, in_channels).

 required
encoder_hidden_states Tensor

Tapped text-encoder hidden states (batch, text_seq_len, num_text_layers, text_hidden_dim).

 required
timestep Tensor

Flow-matching time in [0, 1] of shape (batch,).

 required
position_ids Tensor

(t, h, w) rotary coordinates (text_seq_len + image_seq_len, 3).

 required
encoder_attention_mask Tensor | None

Boolean mask marking valid text tokens (batch, text_seq_len) or None.

 None

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

Krea2TransformerBlock

Bases: Module

attn instance-attribute 

attn = Krea2Attention(
    hidden_size,
    num_heads,
    num_kv_heads,
    eps=norm_eps,
    role="self",
    quant_config=quant_config,
    prefix=_join_prefix(prefix, "attn"),
)

ff instance-attribute 

ff = Krea2SwiGLU(
    hidden_size,
    intermediate_size,
    quant_config=quant_config,
    prefix=_join_prefix(prefix, "ff"),
)

norm1 instance-attribute 

norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps)

norm2 instance-attribute 

norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps)

scale_shift_table instance-attribute 

scale_shift_table = nn.Parameter(
    torch.zeros(6, hidden_size)
)

forward 

forward(
    hidden_states: Tensor,
    temb: Tensor,
    image_rotary_emb: tuple[Tensor, Tensor],
    attention_mask: Tensor | None = None,
) -> Tensor