Skip to content

vllm_omni.diffusion.models.krea2

Krea 2 diffusion model components.

Modules:

Name Description
krea2_transformer
pipeline_krea2

Krea2Pipeline

Bases: Module, DiffusionPipelineProfilerMixin, ProgressBarMixin, SupportsComponentDiscovery

The Krea 2 text-to-image pipeline (vLLM-Omni port of diffusers.Krea2Pipeline).

Components: - scheduler: FlowMatchEulerDiscreteScheduler (resolution-aware exponential time shift). - vae: the Qwen-Image VAE (f8, 16 latent channels). - text_encoder: a Qwen3-VL model; the pipeline consumes a stack of hidden states tapped from several decoder layers rather than the last hidden state. - tokenizer: the tokenizer paired with the text encoder. - transformer: the Krea 2 single-stream MMDiT that predicts the flow-matching velocity.

default_sample_size instance-attribute 

default_sample_size = 1024

device instance-attribute 

device = get_local_device()

do_classifier_free_guidance property 

do_classifier_free_guidance

image_processor instance-attribute 

image_processor = VaeImageProcessor(
    vae_scale_factor=self.vae_scale_factor * self.patch_size
)

is_distilled instance-attribute 

is_distilled = bool(model_index.get('is_distilled', False))

od_config instance-attribute 

od_config = od_config

parallel_config instance-attribute 

parallel_config = od_config.parallel_config

patch_size instance-attribute 

patch_size = int(model_index.get('patch_size', 2))

prompt_template_encode_num_suffix_tokens instance-attribute 

prompt_template_encode_num_suffix_tokens = 5

prompt_template_encode_prefix instance-attribute 

prompt_template_encode_prefix = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"

prompt_template_encode_start_idx instance-attribute 

prompt_template_encode_start_idx = 34

prompt_template_encode_suffix instance-attribute 

prompt_template_encode_suffix = (
    "<|im_end|>\n<|im_start|>assistant\n"
)

scheduler instance-attribute 

scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
    model,
    subfolder="scheduler",
    local_files_only=local_files_only,
)

text_encoder instance-attribute 

text_encoder = self.text_encoder.to(self.device)

text_encoder_select_layers instance-attribute 

text_encoder_select_layers = tuple(
    (int(i)) for i in select_layers
)

tokenizer instance-attribute 

tokenizer = AutoTokenizer.from_pretrained(
    model,
    subfolder="tokenizer",
    local_files_only=local_files_only,
)

transformer instance-attribute 

transformer = Krea2Transformer2DModel(
    od_config=od_config,
    quant_config=od_config.quantization_config,
    **transformer_kwargs,
)

vae instance-attribute 

vae = from_pretrained_with_prefetch(
    DistributedAutoencoderKLQwenImage.from_pretrained,
    model,
    subfolder="vae",
    prefetch_list=subfolders,
    local_files_only=local_files_only,
    torch_dtype=od_config.dtype,
).to(self.device)

vae_scale_factor instance-attribute 

vae_scale_factor = (
    2 ** len(self.vae.temperal_downsample)
    if getattr(self, "vae", None)
    else 8
)

weights_sources instance-attribute 

weights_sources = [
    DiffusersPipelineLoader.ComponentSource(
        model_or_path=od_config.model,
        subfolder="transformer",
        revision=od_config.revision,
        prefix="transformer.",
        fall_back_to_pt=True,
    )
]

diffuse 

diffuse(
    latents: Tensor,
    timesteps: Tensor,
    position_ids: Tensor,
    prompt_embeds: Tensor,
    prompt_embeds_mask: Tensor,
    negative_prompt_embeds: Tensor | None,
    negative_prompt_embeds_mask: Tensor | None,
    guidance_scale: float,
) -> Tensor

Run the flow-matching denoising loop and return the final latents.

encode_prompt 

encode_prompt(
    prompt: str | list[str],
    device: device | None = None,
    num_images_per_prompt: int = 1,
    prompt_embeds: Tensor | None = None,
    prompt_embeds_mask: Tensor | None = None,
    max_sequence_length: int = 512,
) -> tuple[Tensor, Tensor]

forward 

forward(
    req: OmniDiffusionRequest,
    prompt: str | list[str] | None = None,
    negative_prompt: str | list[str] | None = None,
    height: int = 1024,
    width: int = 1024,
    num_inference_steps: int = 28,
    sigmas: list[float] | None = None,
    guidance_scale: float = 4.5,
    num_images_per_prompt: int = 1,
    generator: Generator | list[Generator] | None = None,
    latents: Tensor | None = None,
    output_type: str | None = "pil",
    max_sequence_length: int = 512,
) -> DiffusionOutput

get_text_hidden_states 

get_text_hidden_states(
    prompt: str | list[str],
    max_sequence_length: int = 512,
    device: device | None = None,
) -> tuple[Tensor, Tensor]

Tokenize prompt into the fixed-length Krea 2 layout and tap the selected encoder hidden states.

Returns (hidden_states, attention_mask) of shapes (batch, text_seq_len, num_text_layers, text_hidden_dim) and (batch, text_seq_len) (bool).

load_weights 

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

prepare_latents 

prepare_latents(
    batch_size,
    num_channels_latents,
    height,
    width,
    dtype,
    device,
    generator,
    latents=None,
)

prepare_position_ids staticmethod 

prepare_position_ids(
    text_seq_len: int,
    grid_height: int,
    grid_width: int,
    device: device,
)

Build the (text_seq_len + grid_height * grid_width, 3) rotary coordinates: text tokens sit at the origin, image tokens carry their (0, h, w) latent-grid coordinates.

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]

get_krea2_post_process_func 

get_krea2_post_process_func(od_config: OmniDiffusionConfig)