Diffusers documentation

Krea 2

Diffusers

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.39.0).

Join the Hugging Face community

and get access to the augmented documentation experience

Collaborate on models, datasets and Spaces

Faster examples with accelerated inference

Switch between documentation themes

to get started

Krea 2

Krea 2 (K2) is a flow-matching text-to-image model built around a single-stream MMDiT with grouped-query attention. A Qwen3-VL text encoder provides the conditioning: instead of the last hidden state, hidden states from twelve decoder layers are tapped per token and fused inside the transformer by a small text-fusion stage. Images are decoded with the Qwen-Image VAE.

Two checkpoints are released, sharing the same architecture but with different recommended sampler settings:

Base (midtrain) — use the full sampler with classifier-free guidance: num_inference_steps=28, guidance_scale=4.5.
TDM (distilled) — distilled for few-step sampling, run with num_inference_steps=8 and guidance disabled (guidance_scale=0.0).

guidance_scale follows the Krea 2 convention: the velocity is computed as cond + guidance_scale * (cond - uncond) and guidance is enabled whenever guidance_scale > 0 (this equals the usual CFG formulation with scale 1 + guidance_scale).

Text-to-image

import torch
from diffusers import Krea2Pipeline

# Load from a local directory produced by the Krea 2 conversion (no hub repo yet).
pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Raw", torch_dtype=torch.bfloat16)
pipe.to("cuda")

prompt = "a fox in the snow"
image = pipe(
    prompt,
    height=1024,
    width=1024,
    num_inference_steps=28,
    guidance_scale=4.5,
    generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save("krea2.png")

We additionally provide an example for using Krea2 Turbo :

import torch
from diffusers import Krea2Pipeline

pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", torch_dtype=torch.bfloat16)
pipe.to("cuda")

image = pipe(
    "a fox in the snow",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=0.0,
    generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save("krea2_turbo.png")

Krea2Pipeline

class diffusers.Krea2Pipeline

< source >

( scheduler: FlowMatchEulerDiscreteSchedulervae: AutoencoderKLQwenImagetext_encoder: Qwen3VLModeltokenizer: AutoTokenizertransformer: Krea2Transformer2DModeltext_encoder_select_layers: tuple[int, ...] | list[int] | None = Noneis_distilled: bool = Falsepatch_size: int = 2 )

Parameters

scheduler (FlowMatchEulerDiscreteScheduler) — Euler flow-matching scheduler. The Krea 2 sigma schedule is the resolution-aware exponential time shift, so the scheduler config is expected to set use_dynamic_shifting=True together with the Krea 2 shift parameters (base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=6400).
vae (AutoencoderKLQwenImage) — The Qwen-Image variational auto-encoder (f8, 16 latent channels) used to decode latents to images.
text_encoder (PreTrainedModel) — A Qwen3-VL model (e.g. Qwen3VLModel of Qwen/Qwen3-VL-4B-Instruct). The pipeline consumes a stack of hidden states tapped from several decoder layers rather than the last hidden state.
tokenizer (AutoTokenizer) — The tokenizer paired with the text encoder.
transformer (Krea2Transformer2DModel) — The Krea 2 single-stream MMDiT that predicts the flow-matching velocity.
text_encoder_select_layers (tuple[int, ...], optional) — Indices into the text encoder’s hidden_states tuple (0 is the embedding output) whose states are stacked per token as the transformer’s text conditioning. Must have transformer.config.num_text_layers entries.
is_distilled (bool, optional, defaults to False) — Whether the transformer is the few-step distilled (TDM/turbo) checkpoint. When True a fixed timestep shift mu=1.15 is used; otherwise mu is computed from the image resolution.
patch_size (int, optional, defaults to 2) — Side length of the square patches the latents are packed into before entering the transformer. The effective pixel-to-token downsampling factor is vae_scale_factor * patch_size.

The Krea 2 pipeline for text-to-image generation.

call

< source >

( prompt: str | list[str] | None = Nonenegative_prompt: str | list[str] | None = Noneheight: int = 1024width: int = 1024num_inference_steps: int = 28sigmas: list[float] | None = Noneguidance_scale: float = 4.5num_images_per_prompt: int = 1generator: typing.Union[torch.Generator, list[torch.Generator], NoneType] = Nonelatents: typing.Optional[torch.Tensor] = Noneprompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_embeds_mask: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds: typing.Optional[torch.Tensor] = Nonenegative_prompt_embeds_mask: typing.Optional[torch.Tensor] = Noneoutput_type: str | None = 'pil'return_dict: bool = Truecallback_on_step_end: typing.Optional[typing.Callable[[int, int, dict], NoneType]] = Nonecallback_on_step_end_tensor_inputs: list = ['latents']attention_kwargs: dict[str, typing.Any] | None = Nonemax_sequence_length: int = 512 ) → Krea2PipelineOutput or tuple

Parameters

prompt (str or list[str], optional) — The prompt or prompts to guide the image generation. If not defined, one has to pass prompt_embeds.
negative_prompt (str or list[str], optional) — The prompt or prompts not to guide the image generation. Ignored when guidance_scale <= 0; defaults to an empty prompt when guidance is enabled.
height (int, defaults to 1024) — The height in pixels of the generated image. Rounded up to a multiple of 16 if needed.
width (int, defaults to 1024) — The width in pixels of the generated image. Rounded up to a multiple of 16 if needed.
num_inference_steps (int, defaults to 28) — The number of denoising steps. Use 28 for the base (midtrain) checkpoint and 8 for the few-step distilled (TDM) checkpoint.
sigmas (list[float], optional) — Custom sigmas for the scheduler. If not defined, the default linspace(1.0, 1/num_inference_steps, num_inference_steps) grid is used (the resolution-aware shift is applied inside the scheduler).
guidance_scale (float, defaults to 4.5) — Classifier-free guidance scale, following the Krea 2 convention: the velocity is computed as cond + guidance_scale * (cond - uncond) and guidance is enabled whenever guidance_scale > 0 (this equals the usual CFG formulation with scale 1 + guidance_scale). Set to 0.0 to disable (e.g. for the TDM checkpoint).
num_images_per_prompt (int, defaults to 1) — The number of images to generate per prompt.
generator (torch.Generator or list[torch.Generator], optional) — One or more torch generator(s) to make generation deterministic.
latents (torch.Tensor, optional) — Pre-generated noisy latents in packed form (batch_size, image_seq_len, in_channels), sampled from a Gaussian distribution, to be used as inputs for image generation.
prompt_embeds (torch.Tensor, optional) — Pre-generated text embeddings of shape (batch_size, text_seq_len, num_text_layers, text_hidden_dim). If not provided, embeddings are generated from prompt.
prompt_embeds_mask (torch.Tensor, optional) — Boolean mask for prompt_embeds; required when prompt_embeds is passed.
negative_prompt_embeds (torch.Tensor, optional) — Pre-generated negative text embeddings; same layout as prompt_embeds.
negative_prompt_embeds_mask (torch.Tensor, optional) — Boolean mask for negative_prompt_embeds; required when negative_prompt_embeds is passed.
output_type (str, optional, defaults to "pil") — The output format of the generated image. Choose between "pil", "np", "pt" or "latent".
return_dict (bool, optional, defaults to True) — Whether or not to return a Krea2PipelineOutput instead of a plain tuple.
callback_on_step_end (Callable, optional) — A function that is called at the end of each denoising step with callback_on_step_end(self, step, timestep, callback_kwargs).
callback_on_step_end_tensor_inputs (list[str], optional, defaults to ["latents"]) — The list of tensor inputs for the callback_on_step_end function. Must be a subset of ._callback_tensor_inputs.
attention_kwargs (dict, optional) — A kwargs dictionary that if specified is passed along to the AttentionProcessor as defined under self.processor in diffusers.models.attention_processor.
max_sequence_length (int, defaults to 512) — Fixed text sequence length consumed by the transformer; prompts are padded or truncated to it.

Returns

Krea2PipelineOutput or tuple

Krea2PipelineOutput if return_dict is True, otherwise a tuple, whose first element is a list with the generated images.

Function invoked when calling the pipeline for generation.

Examples:

>>> import torch
>>> from diffusers import Krea2Pipeline

>>> # Load from a local directory produced by the Krea 2 conversion (no hub repo yet).
>>> pipe = Krea2Pipeline.from_pretrained("path/to/krea2-diffusers", torch_dtype=torch.bfloat16)
>>> pipe.to("cuda")
>>> prompt = "a fox in the snow"
>>> # Base (midtrain) checkpoint defaults. For the few-step distilled (TDM) checkpoint use
>>> # `num_inference_steps=8, guidance_scale=0.0` instead.
>>> image = pipe(prompt, num_inference_steps=28, guidance_scale=4.5).images[0]
>>> image.save("krea2.png")

encode_prompt

< source >

( prompt: str | list[str]device: typing.Optional[torch.device] = Nonenum_images_per_prompt: int = 1prompt_embeds: typing.Optional[torch.Tensor] = Noneprompt_embeds_mask: typing.Optional[torch.Tensor] = Nonemax_sequence_length: int = 512 )

Parameters

prompt (str or list[str], optional) — prompt to be encoded
device — (torch.device): torch device
num_images_per_prompt (int) — number of images that should be generated per prompt
prompt_embeds (torch.Tensor, optional) — Pre-generated text embeddings of shape (batch_size, text_seq_len, num_text_layers, text_hidden_dim). Can be used to easily tweak text inputs, e.g. prompt weighting. If not provided, text embeddings will be generated from prompt input argument.
prompt_embeds_mask (torch.Tensor, optional) — Pre-generated boolean mask marking valid text tokens, of shape (batch_size, text_seq_len). Required when prompt_embeds is passed.
max_sequence_length (int, defaults to 512) — Fixed text sequence length consumed by the transformer; prompts are padded or truncated to it.

get_text_hidden_states

< source >

( prompt: str | list[str]max_sequence_length: int = 512device: typing.Optional[torch.device] = None )

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

Returns a (hidden_states, attention_mask) tuple of shapes (batch_size, text_seq_len, num_text_layers, text_hidden_dim) and (batch_size, text_seq_len) (bool).

prepare_position_ids

< source >

( text_seq_len: intgrid_height: intgrid_width: intdevice: device )

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

Krea2PipelineOutput

class diffusers.pipelines.krea2.Krea2PipelineOutput

< source >

( images: list[PIL.Image.Image] | numpy.ndarray )

Parameters

images (list[PIL.Image.Image] or np.ndarray) — List of denoised PIL images of length batch_size or numpy array of shape (batch_size, height, width, num_channels).

Output class for the Krea 2 pipeline.

Modular

Krea 2 is also available as a modular pipeline. Classifier-free guidance is configured through the guider component rather than a guidance_scale call argument. Krea 2 uses cond-anchored CFG, which is ClassifierFreeGuidance with use_original_formulation=True.

import torch
from diffusers import ClassifierFreeGuidance, ModularPipeline

pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw")
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")


image = pipe(
    prompt="a fox in the snow",
    height=1024,
    width=1024,
    num_inference_steps=28,
    generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save("krea2.png")

We additionally provide an example for using Krea2 Turbo. The distilled checkpoint maps to its own set of blocks (Krea2TurboAutoBlocks): it runs guidance-free (no guider), takes no negative prompt, and samples in a few steps. ModularPipeline.from_pretrained picks the turbo blocks automatically from the checkpoint’s is_distilled config, so no guidance configuration is needed:

import torch
from diffusers import ModularPipeline

pipe = ModularPipeline.from_pretrained("krea/Krea-2-Turbo")
pipe.load_components(torch_dtype=torch.bfloat16)
pipe.to("cuda")

image = pipe(
    prompt="a fox in the snow",
    height=1024,
    width=1024,
    num_inference_steps=8,
    generator=torch.Generator("cuda").manual_seed(0),
).images[0]
image.save("krea2_turbo.png")

Krea2ModularPipeline

class diffusers.Krea2ModularPipeline

< source >

( blocks: diffusers.modular_pipelines.modular_pipeline.ModularPipelineBlocks | None = Nonepretrained_model_name_or_path: str | os.PathLike | None = Nonecomponents_manager: diffusers.modular_pipelines.components_manager.ComponentsManager | None = Nonecollection: str | None = Nonemodular_config_dict: dict[str, typing.Any] | None = Noneconfig_dict: dict[str, typing.Any] | None = None**kwargs )

A ModularPipeline for Krea 2.

> This is an experimental feature!

Krea2AutoBlocks

class diffusers.Krea2AutoBlocks

< source >

( )

Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise (symmetric CFG) -> decode.

Supported workflows:

text2image: requires prompt

Components: text_encoder (Qwen3VLModel): The Qwen3-VL text encoder. tokenizer (AutoTokenizer): The tokenizer paired with the text encoder. guider (ClassifierFreeGuidance) transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Inputs: prompt (str): The prompt or prompts to guide image generation. negative_prompt (str, optional): The negative prompt(s) for CFG. max_sequence_length (int, optional, defaults to 512): Maximum sequence length for prompt encoding. num_images_per_prompt (int, optional, defaults to 1): The number of images to generate per prompt. latents (Tensor, optional): Pre-generated noisy latents for image generation. height (int, optional, defaults to 1024): The height in pixels of the generated image. width (int, optional, defaults to 1024): The width in pixels of the generated image. generator (Generator, optional): Torch generator for deterministic generation. num_inference_steps (int, optional, defaults to 28): The number of denoising steps. sigmas (list, optional): Custom sigma schedule (defaults to a linear ramp). attention_kwargs (dict, optional): Additional kwargs for attention processors. output_type (str, optional, defaults to pil): Output format: ‘pil’, ‘np’, ‘pt’.

Outputs: images (list): Generated images.

Krea2TurboModularPipeline

class diffusers.Krea2TurboModularPipeline

< source >

A ModularPipeline for the distilled Krea 2 turbo (TDM) checkpoint. It runs without classifier-free guidance, so it takes no negative prompt and has no guider.

> This is an experimental feature!

Krea2TurboAutoBlocks

class diffusers.Krea2TurboAutoBlocks

< source >

( )

Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode text -> core denoise (guidance-free) -> decode.

Supported workflows:

text2image: requires prompt

Components: text_encoder (Qwen3VLModel): The Qwen3-VL text encoder. tokenizer (AutoTokenizer): The tokenizer paired with the text encoder. transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Inputs: prompt (str): The prompt or prompts to guide image generation. max_sequence_length (int, optional, defaults to 512): Maximum sequence length for prompt encoding. num_images_per_prompt (int, optional, defaults to 1): The number of images to generate per prompt. latents (Tensor, optional): Pre-generated noisy latents for image generation. height (int, optional, defaults to 1024): The height in pixels of the generated image. width (int, optional, defaults to 1024): The width in pixels of the generated image. generator (Generator, optional): Torch generator for deterministic generation. num_inference_steps (int, optional, defaults to 8): The number of denoising steps. sigmas (list, optional): Custom sigma schedule (defaults to a linear ramp). attention_kwargs (dict, optional): Additional kwargs for attention processors. output_type (str, optional, defaults to pil): Output format: ‘pil’, ‘np’, ‘pt’.

Outputs: images (list): Generated images.

Update on GitHub

←Kolors Latent Consistency Models→

Diffusers

Krea 2

Text-to-image

Krea2Pipeline

class diffusers.Krea2Pipeline

__call__

encode_prompt

get_text_hidden_states

prepare_position_ids

Krea2PipelineOutput

class diffusers.pipelines.krea2.Krea2PipelineOutput

Modular

Krea2ModularPipeline

class diffusers.Krea2ModularPipeline

Krea2AutoBlocks

class diffusers.Krea2AutoBlocks

Krea2TurboModularPipeline

class diffusers.Krea2TurboModularPipeline

Krea2TurboAutoBlocks

class diffusers.Krea2TurboAutoBlocks

call