Skip to content

vllm_omni.diffusion.models.krea2.pipeline_krea2

DEFAULT_TEXT_ENCODER_SELECT_LAYERS module-attribute 

DEFAULT_TEXT_ENCODER_SELECT_LAYERS = (
    2,
    5,
    8,
    11,
    14,
    17,
    20,
    23,
    26,
    29,
    32,
    35,
)

logger module-attribute 

logger = logging.getLogger(__name__)

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.

calculate_shift 

calculate_shift(
    image_seq_len,
    base_seq_len: int = 256,
    max_seq_len: int = 4096,
    base_shift: float = 0.5,
    max_shift: float = 1.15,
)

get_krea2_post_process_func 

get_krea2_post_process_func(od_config: OmniDiffusionConfig)

retrieve_timesteps 

retrieve_timesteps(
    scheduler,
    num_inference_steps: int | None = None,
    device: str | device | None = None,
    timesteps: list[int] | None = None,
    sigmas: list[float] | None = None,
    **kwargs,
) -> tuple[Tensor, int]