Diffusers documentation
Krea 2
Krea 2
Krea 2 (K2) 是一个基于单流 MMDiT 和分组查询注意力机制的流匹配文本生成图像模型。Qwen3-VL 文本编码器提供条件信息：对于每个 token，不再使用最后一个隐藏状态，而是从十二个解码器层中提取隐藏状态，并通过 transformer 内部的一个小型文本融合阶段进行融合。图像通过 Qwen-Image VAE 进行解码。
发布了两个检查点，它们共享相同的架构，但具有不同的推荐采样器设置：
- 基础（中期训练） — 使用完整的采样器并配合无分类器引导：
num_inference_steps=28，guidance_scale=4.5。 - TDM（蒸馏版） — 专为少步采样而蒸馏，运行时使用
num_inference_steps=8并禁用引导（guidance_scale=0.0）。
guidance_scale 遵循 Krea 2 的约定：速度计算为 cond + guidance_scale * (cond - uncond)，并且只要 guidance_scale > 0 就会启用 guidance（这等同于使用 scale 1 + guidance_scale 的常规 CFG 公式）。
文本到图像
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", dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
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")我们还提供了一个使用 Krea2 Turbo 的示例：
import torch
from diffusers import Krea2Pipeline
pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
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")加载单文件检查点
import torch
from diffusers import Krea2Pipeline, Krea2Transformer2DModel
transformer = Krea2Transformer2DModel.from_single_file(
"https://huggingface.co/krea/Krea-2-Turbo/blob/main/turbo.safetensors", dtype=torch.bfloat16
)
pipe = Krea2Pipeline.from_pretrained("krea/Krea-2-Turbo", transformer=transformer, dtype=torch.bfloat16).to("cuda")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 )
参数
- scheduler (FlowMatchEulerDiscreteScheduler) —
Euler 流匹配调度器。Krea 2 的 sigma 调度采用分辨率感知的指数时间偏移，因此调度器配置预期会设置
use_dynamic_shifting=True，并配合 Krea 2 的偏移参数（base_shift=0.5、max_shift=1.15、base_image_seq_len=256、max_image_seq_len=6400）。 - vae (AutoencoderKLQwenImage) — 用于将潜在变量解码为图像的 Qwen-Image 变分自编码器（f8，16 个潜在通道）。
- text_encoder (PreTrainedModel) —
一个 Qwen3-VL 模型（例如
Qwen/Qwen3-VL-4B-Instruct的Qwen3VLModel）。该 pipeline 使用从多个解码器层中提取的隐藏状态堆栈，而不是最后一个隐藏状态。 - tokenizer（AutoTokenizer）—— 与文本编码器配对的 tokenizer。
- transformer (Krea2Transformer2DModel) — Krea 2 单流 MMDiT，用于预测流匹配速度。
- text_encoder_select_layers (
tuple[int, ...]，可选) — 文本编码器hidden_states元组（0 为嵌入输出）中，其状态将按 token 堆叠并作为 transformer 文本条件的索引。必须包含transformer.config.num_text_layers个条目。 - is_distilled (
bool，可选，默认为False) — 指示该 transformer 是否为少步蒸馏（TDM/turbo）检查点。当True时，使用固定的时间步偏移mu=1.15；否则mu将根据图像分辨率计算得出。 - patch_size (
int，可选，默认为 2) — 在潜变量进入 transformer 之前，将其打包成的方形 patch 的边长。有效的像素到 token 的下采样因子为vae_scale_factor * patch_size。
用于文本到图像生成的 Krea 2 流水线。
__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
参数
- prompt (
str或list[str]，可选) — 用于引导图像生成的提示词或提示词列表。如果未定义，则必须传入prompt_embeds。 - negative_prompt (
str或list[str]，可选) — 用于引导图像生成的反向提示词。当guidance_scale <= 0时将被忽略；在启用引导时，默认值为空提示词。 - height (
int，默认为 1024) — 生成图像的高度（以像素为单位）。如有必要，会向上取整为 16 的倍数。 - width (
int，默认为 1024) — 生成图像的宽度（以像素为单位）。如有必要，会向上取整为 16 的倍数。 - num_inference_steps (
int，默认为 28) — 去噪步数。对于基础（midtrain）检查点，请使用 28；对于少步蒸馏（TDM）检查点，请使用 8。 - sigmas (
list[float]，可选) — 调度器的自定义 sigmas。如果未定义，则使用默认的linspace(1.0, 1/num_inference_steps, num_inference_steps)网格（分辨率感知的 shift 在调度器内部应用）。 - guidance_scale (
float，默认为 4.5) — 无分类器引导（Classifier-free guidance）的缩放系数，遵循 Krea 2 的约定：速度计算为cond + guidance_scale * (cond - uncond)，且只要guidance_scale > 0即启用引导（这等同于使用缩放系数1 + guidance_scale的常规 CFG 公式）。将其设置为0.0可禁用引导（例如用于 TDM 检查点）。 - num_images_per_prompt (
int，默认为 1) — 每个提示词生成的图像数量。 - generator (
torch.Generator或list[torch.Generator]，可选) — 一个或多个 torch 生成器，用于使生成过程具有确定性。 - latents (
torch.Tensor, 可选) — 以打包形式预生成的含噪潜变量(batch_size, image_seq_len, in_channels)，从 高斯分布中采样，用作图像生成的输入。 - prompt_embeds (
torch.Tensor, 可选) — 预生成的文本嵌入，形状为(batch_size, text_seq_len, num_text_layers, text_hidden_dim)。 如果未提供，则从prompt生成嵌入。 - prompt_embeds_mask (
torch.Tensor, 可选) —prompt_embeds的布尔掩码；当传入prompt_embeds时必填。 - negative_prompt_embeds (
torch.Tensor，可选) — 预生成的负向文本嵌入；布局与prompt_embeds相同。 - negative_prompt_embeds_mask (
torch.Tensor, 可选) —negative_prompt_embeds的布尔掩码；当传入negative_prompt_embeds时必填。 - output_type (
str，可选，默认为"pil") — 生成图像的输出格式。可选"pil"、"np"、"pt"或"latent"。 - return_dict (
bool，可选，默认为True) — 是否返回 Krea2PipelineOutput 而不是普通元组。 - callback_on_step_end (
Callable, 可选) — 在每个去噪步骤结束时调用的函数，调用形式为callback_on_step_end(self, step, timestep, callback_kwargs)。 - callback_on_step_end_tensor_inputs (
list[str]，可选，默认为["latents"]) —callback_on_step_end函数的张量输入列表。必须是._callback_tensor_inputs的子集。 - attention_kwargs (
dict, 可选) — 一个 kwargs 字典，如果指定，将传递给 diffusers.models.attention_processor 中定义的AttentionProcessor，该处理器在self.processor下定义。 - max_sequence_length (
int，默认为 512) — Transformer 消耗的固定文本序列长度；提示词会被填充或截断至该长度。
返回值
Krea2PipelineOutput 或 tuple
如果 return_dict 为 True，则返回 Krea2PipelineOutput；否则返回一个 tuple，其第一个元素是包含生成图像的列表。
调用生成 pipeline 时调用的函数。
示例：
>>> 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 )
参数
- prompt (
str或list[str]，可选) — 要编码的提示 - device — (
torch.device)： torch 设备 - num_images_per_prompt (
int) — 每个提示词应生成的图像数量 - prompt_embeds (
torch.Tensor, 可选) — 预生成的文本嵌入，形状为(batch_size, text_seq_len, num_text_layers, text_hidden_dim)。 可用于轻松调整文本输入，例如 提示词加权。如果未提供，文本嵌入将从prompt输入参数生成。 - prompt_embeds_mask (
torch.Tensor, 可选) — 预生成的布尔掩码，用于标记有效的文本 token，形状为(batch_size, text_seq_len)。当传入prompt_embeds时必填。 - max_sequence_length (
int，默认为 512) — Transformer 消耗的固定文本序列长度；提示词会被填充或截断至该长度。
( prompt: str | list[str]max_sequence_length: int = 512device: typing.Optional[torch.device] = None )
将 prompt 分词为固定长度的 Krea 2 布局，并获取所选编码器的隐藏状态。
返回一个 (hidden_states, attention_mask) 元组，其形状分别为 (batch_size, text_seq_len, num_text_layers, text_hidden_dim) 和 (batch_size, text_seq_len)（bool 类型）。
构建合并序列的 (text_seq_len + grid_height * grid_width, 3) 旋转坐标：
文本 token 位于原点，图像 token 携带其 (0, h, w) 潜在网格坐标。
Krea2PipelineOutput
类 diffusers.pipelines.krea2.Krea2PipelineOutput
< source >( images: list[PIL.Image.Image] | numpy.ndarray )
Krea 2 流水线的输出类。
模块化
Krea 2 也可作为模块化流水线使用。无分类器引导通过guider组件进行配置，而不是通过guidance_scale调用参数。Krea 2 使用基于条件锚定的 CFG，即带有use_original_formulation=True的ClassifierFreeGuidance。
import torch
from diffusers import ClassifierFreeGuidance, ModularPipeline
pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw")
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
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")我们额外提供了一个使用 Krea2 Turbo 的示例。蒸馏后的检查点映射到其专属的一组块（Krea2TurboAutoBlocks）：它以无引导方式运行（不使用 guider），不接受负面提示词，并且只需少量步骤即可完成采样。ModularPipeline.from_pretrained 会根据检查点的 is_distilled 配置自动选择 turbo 块，因此
无需配置引导：
import torch
from diffusers import ModularPipeline
pipe = ModularPipeline.from_pretrained("krea/Krea-2-Turbo")
pipe.load_components(dtype=torch.bfloat16)
pipe.to("cuda") # or "mps", "xpu", "cpu"
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 = Noneworkflow: str | None = Nonemodular_config_dict: dict[str, typing.Any] | None = Noneconfig_dict: dict[str, typing.Any] | None = None**kwargs )
用于 Krea 2 的 ModularPipeline。
Krea2AutoBlocks
使用 Krea 2 进行文本到图像生成的 Auto Modular 流水线：编码文本 -> 核心去噪（对称 CFG） -> 解码。
支持的工作流：
text2image：需要prompt
组件：
text_encoder (Qwen3VLModel)：Qwen3-VL 文本编码器。 tokenizer (AutoTokenizer)：与文本编码器配对的 tokenizer。 guider (ClassifierFreeGuidance) transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)
输入：
prompt (str)：
用于引导图像生成的提示词或提示词列表。
negative_prompt (str，可选)：
用于 CFG 的负向提示词。
max_sequence_length (int，可选，默认为 512)：
提示词编码的最大序列长度。
num_images_per_prompt (int，可选，默认为 1)：
每个提示词生成的图像数量。
latents (Tensor，可选)：
用于图像生成的预生成噪声潜变量。
height (int，可选，默认为 1024)：
生成图像的高度（以像素为单位）。
width (int，可选，默认为 1024)：
生成图像的宽度（以像素为单位）。
generator (Generator，可选)：
用于确定性生成的 Torch 生成器。
num_inference_steps (int，可选，默认为 28)：
去噪步数。
sigmas (list，可选)：
自定义 sigma 调度（默认为线性斜坡）。
attention_kwargs (dict，可选)：
注意力处理器的额外关键字参数。
output_type (str，可选，默认为 pil)：
输出格式：‘pil’、‘np’、‘pt’。
输出：
images (list)：
生成的图像。
Krea2TurboModularPipeline
类 diffusers.Krea2TurboModularPipeline
< 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 = Noneworkflow: str | None = Nonemodular_config_dict: dict[str, typing.Any] | None = Noneconfig_dict: dict[str, typing.Any] | None = None**kwargs )
用于蒸馏版 Krea 2 turbo (TDM) 检查点的 ModularPipeline。它在运行时不使用无分类器引导，因此不接受负面提示，也没有引导器。
Krea2TurboAutoBlocks
使用蒸馏版 Krea 2 turbo 检查点实现文本到图像生成的 Auto Modular 流水线：文本编码 -> 核心去噪（无引导） -> 解码。
支持的工作流：
text2image：需要prompt
组件：
text_encoder (Qwen3VLModel)：Qwen3-VL 文本编码器。 tokenizer (AutoTokenizer)：与文本编码器配套的 tokenizer。 transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler)
vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)
输入：
prompt (str)：
用于引导图像生成的提示词或提示词列表。
max_sequence_length (int，可选，默认为 512)：
提示词编码的最大序列长度。
num_images_per_prompt (int，可选，默认为 1)：
每个提示词生成的图像数量。
latents (Tensor，可选)：
用于图像生成的预生成含噪潜在变量。
height (int，可选，默认为 1024)：
生成图像的高度（以像素为单位）。
width (int，可选，默认为 1024)：
生成图像的宽度（以像素为单位）。
generator (Generator，可选)：
用于确定性生成的 Torch 生成器。
num_inference_steps (int，可选，默认为 8)：
去噪步数。
sigmas (list，可选)：
自定义 sigma 调度（默认为线性斜坡）。
attention_kwargs (dict，可选)：
注意力处理器的额外关键字参数。
output_type (str，可选，默认为 pil)：
输出格式：‘pil’、‘np’、‘pt’。
输出：
images (list)：
生成的图像。