Diffusers documentation

Krea 2

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v0.42.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

Krea 2

Krea 2 (K2) は、グループクエリアテンションを備えた単一ストリームの MMDiT を基盤とするフローマッチング型のテキストから画像を生成するモデルです。条件付けには Qwen3-VL テキストエンコーダーが使用されます。最終層の隠れ状態ではなく、トークンごとに 12 個のデコーダー層の隠れ状態を取得し、トランスフォーマー内部の小型のテキスト融合ステージで融合します。画像のデコードには Qwen-Image VAE を使用します。

2つのチェックポイントが公開されており、同じアーキテクチャを共有していますが、推奨されるサンプラーの設定が異なります。

  • Base (midtrain) — classifier-free guidance を使用したフルサンプラーを使用します: 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 が有効になります(これは、スケールが 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

クラス diffusers.Krea2Pipeline

< >

( 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 flow-matching スケジューラ。Krea 2 の sigma スケジュールは解像度認識型の指数時間シフトであるため、スケジューラの設定では Krea 2 のシフトパラメータ (base_shift=0.5、max_shift=1.15、base_image_seq_len=256、max_image_seq_len=6400) とともに use_dynamic_shifting=True が設定されることを想定しています。
  • vae (AutoencoderKLQwenImage) — Qwen-Image の変分オートエンコーダー(f8、16 個の潜在チャネル)です。潜在変数から画像をデコードするために使用されます。
  • text_encoder (PreTrainedModel) — Qwen3-VL モデル(例: Qwen/Qwen3-VL-4B-Instruct の Qwen3VLModel)です。このパイプラインは、最後の hidden state ではなく、複数のデコーダー層から取得した hidden states のスタックを消費します。
  • tokenizer (AutoTokenizer) — テキストエンコーダーと組み合わせて使用するトークナイザーです。
  • transformer (Krea2Transformer2DModel) — フローマッチングの速度を予測する Krea 2 のシングルストリーム MMDiT です。
  • text_encoder_select_layers (tuple[int, ...]、省略可) — テキストエンコーダーのhidden_statesタプル(0 は埋め込み出力)内のインデックス。これらの状態がトークンごとにスタックされ、transformer のテキスト条件付けとして使用されます。transformer.config.num_text_layers 個のエントリを持つ必要があります。
  • is_distilled (bool、任意、デフォルトは False) — transformer がステップ数の少ない蒸留済み(TDM/turbo)チェックポイントかどうかを示します。True の場合、固定のタイムステップシフト mu=1.15 が使用されます。それ以外の場合は、mu が画像の解像度から計算されます。
  • patch_size (int, 任意、デフォルトは 2) — transformer に渡す前に、latents をパッキングする正方形パッチの辺の長さです。 ピクセルからトークンへの実効的なダウンサンプリング係数は vae_scale_factor * patch_size です。

テキストから画像を生成するための Krea 2 パイプライン。

__call__

< >

( 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[str] = ['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]、任意) — スケジューラ用のカスタムシグマ。未定義の場合、デフォルトのlinspace(1.0, 1/num_inference_steps, num_inference_steps)グリッドが使用されます(解像度認識シフトはスケジューラ内で適用されます)。
  • guidance_scale (float、デフォルト値は 4.5) — Krea 2 の慣例に従う Classifier-free guidance のスケールです。velocity は cond + guidance_scale * (cond - uncond) として計算され、guidance_scale > 0 の場合常に guidance が有効になります(これは、スケールが 1 + guidance_scale である通常の CFG 定式化に相当します)。無効にするには 0.0 を設定してください(例:TDM チェックポイントの場合)。
  • num_images_per_prompt (int、デフォルト値は 1) — プロンプトごとに生成する画像の数。
  • generator (torch.Generator または list[torch.Generator]、省略可) — 生成を決定論的にするために使用する、1 つ以上の 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) — 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, 省略可) — 指定された場合、diffusers.models.attention_processor で定義されている self.processor 内の AttentionProcessor に渡される kwargs 辞書です。
  • max_sequence_length (int、デフォルト値は 512) — transformer が消費する固定のテキストシーケンス長です。プロンプトは、この長さにパディングまたは切り詰められます。

戻り値

Krea2PipelineOutput または tuple

return_dict が True の場合は Krea2PipelineOutput、それ以外の場合は tuple が返され、その最初の要素は生成された画像のリストです。

生成用のパイプラインを呼び出したときに実行される関数。

例:

>>> 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

< >

( 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) — 1つのプロンプトに対して生成する画像の数
  • prompt_embeds (torch.Tensor, 任意) — (batch_size, text_seq_len, num_text_layers, text_hidden_dim) の形状を持つ、事前生成されたテキスト埋め込みです。 プロンプトの重み付けなど、例えばテキスト入力の調整を簡単に実行するために使用できます。指定されない場合、prompt 入力引数からテキスト埋め込みが生成されます。
  • prompt_embeds_mask (torch.Tensor, 任意) — 有効なテキストトークンを示す事前生成済みのブールマスクで、形状は (batch_size, text_seq_len) です。prompt_embeds が渡される場合に必要です。
  • max_sequence_length (int、デフォルト値は 512) — transformer が消費する固定のテキストシーケンス長です。プロンプトは、この長さにパディングまたは切り詰められます。

get_text_hidden_states

< >

( 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) です。

prepare_position_ids

< >

( text_seq_len: intgrid_height: intgrid_width: intdevice: device )

結合されたシーケンス用の (text_seq_len + grid_height * grid_width, 3) のロータリー座標を構築します: テキストトークンは原点に配置され、画像トークンはその (0, h, w) の潜在グリッド座標を持ちます。

Krea2PipelineOutput

class diffusers.pipelines.krea2.Krea2PipelineOutput

< >

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

パラメータ

  • images (list[PIL.Image.Image] または np.ndarray) — batch_size 長さのデノイズ済み PIL 画像のリスト、または (batch_size, height, width, num_channels) の形状を持つ numpy 配列。

Krea 2 パイプラインの出力クラスです。

モジュール化

Krea 2 は モジュラーパイプライン としても利用できます。Classifier-free guidance は guidance_scale の呼び出し引数ではなく、guider コンポーネントを通じて設定します。Krea 2 は cond-anchored 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) にマッピングされます。guidance なし (guider なし) で動作し、negative prompt を受け取らず、数ステップでサンプリングを行います。チェックポイントの is_distilled 設定から ModularPipeline.from_pretrained が自動的に turbo ブロックを選択するため、guidance の設定は不要です:

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

< >

( 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

class diffusers.Krea2AutoBlocks

< >

( )

Krea 2 を使用したテキストから画像への生成のための Auto Modular パイプライン: テキストのエンコード -> コアデノイズ(対称 CFG) -> デコード。

サポートされているワークフロー:

  • text2image: prompt が必要です

コンポーネント: text_encoder (Qwen3VLModel): Qwen3-VL のテキストエンコーダー。 tokenizer (AutoTokenizer): テキストエンコーダーと組み合わせて使用するトークナイザー。 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): 1つのプロンプトあたりに生成する画像の数。 latents (Tensor, 任意): 画像生成用の事前生成されたノイズ付き潜在変数。 height (int, 任意, デフォルト値: 1024): 生成される画像の高さ(ピクセル単位)。 width (int, 任意, デフォルト値: 1024): 生成される画像の幅(ピクセル単位)。 generator (Generator, 任意): 決定論的な生成のためのTorchジェネレーター。 num_inference_steps (int, 任意, デフォルト値: 28): ノイズ除去ステップの数。 sigmas (list, 任意): カスタムのシグマスケジュール(デフォルトは線形ランプ)。 attention_kwargs (dict, 任意): アテンションプロセッサー用の追加のkwargs。 output_type (str, 任意, デフォルト値: pil): 出力形式: ‘pil’、‘np’、‘pt’。

出力: images (list): 生成された画像。

Krea2TurboModularPipeline

class diffusers.Krea2TurboModularPipeline

< >

( 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 です。classifier-free guidance なしで動作するため、negative prompt を受け取らず、guider も持ちません。

Krea2TurboAutoBlocks

class diffusers.Krea2TurboAutoBlocks

< >

( )

蒸留された Krea 2 turbo チェックポイントを使用してテキストから画像を生成するための Auto Modular パイプライン: テキストのエンコード -> コアデノイズ (ガイダンスなし) -> デコード。

サポートされているワークフロー:

  • text2image: prompt が必要です

コンポーネント: text_encoder (Qwen3VLModel): Qwen3-VL のテキストエンコーダー。 tokenizer (AutoTokenizer): テキストエンコーダーと組み合わせて使用するトークナイザー。 transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

入力: prompt (str): 画像生成をガイドするためのプロンプトまたはプロンプトのリスト。 max_sequence_length (int, 省略可能、デフォルトは 512): プロンプトエンコードの最大シーケンス長。 num_images_per_prompt (int, 省略可能、デフォルトは 1): 1 つのプロンプトあたりに生成する画像の数。 latents (Tensor, 省略可能): 画像生成用の事前生成されたノイズ付き潜在変数。 height (int, 省略可能、デフォルトは 1024): 生成される画像の高さ(ピクセル単位)。 width (int, 省略可能、デフォルトは 1024): 生成される画像の幅(ピクセル単位)。 generator (Generator, 省略可能): 決定論的な生成のための Torch ジェネレーター。 num_inference_steps (int, 省略可能、デフォルトは 8): ノイズ除去ステップの数。 sigmas (list, 省略可能): カスタムのシグマスケジュール(デフォルトは線形ランプ)。 attention_kwargs (dict, 省略可能): アテンションプロセッサ用の追加の kwargs。 output_type (str, 省略可能、デフォルトは pil): 出力形式: ‘pil’、‘np’、‘pt’。

出力: images (list): 生成された画像。

Update on GitHub