Diffusers documentation

Krea 2

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) est un modèle de génération d'images à partir de texte par flow-matching, construit autour d'un MMDiT à flux unique avec une attention à requêtes groupées. Un encodeur de texte Qwen3-VL fournit le conditionnement : au lieu de l'état caché final, les états cachés de douze couches décodeur sont extraits par jeton et fusionnés à l'intérieur du transformeur par un petit stade de fusion de texte. Les images sont décodées avec le VAE Qwen-Image.

Deux points de contrôle sont publiés, partageant la même architecture mais avec des paramètres de sampler recommandés différents :

  • Base (midtrain) — utilisez le sampler complet avec classifier-free guidance : num_inference_steps=28, guidance_scale=4.5.
  • TDM (distillé) — distillé pour un échantillonnage en peu d'étapes, à exécuter avec num_inference_steps=8 et le guidage désactivé (guidance_scale=0.0).

guidance_scale suit la convention de Krea 2 : la vitesse est calculée comme cond + guidance_scale * (cond - uncond) et le guidage est activé dès que guidance_scale > 0 (ce qui équivaut à la formulation habituelle du CFG avec une échelle de 1 + guidance_scale).

Texte vers 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", 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")

Nous fournissons également un exemple d'utilisation de 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")

Chargement de points de contrôle en fichier unique

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

classe diffusers.Krea2Pipeline

< >

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

Paramètres

  • scheduler (FlowMatchEulerDiscreteScheduler) — Planificateur de correspondance de flux d'Euler. Le plan de sigma de Krea 2 est le décalage temporel exponentiel tenant compte de la résolution, la configuration du planificateur est donc censée définir use_dynamic_shifting=True avec les paramètres de décalage de Krea 2 (base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=6400).
  • vae (AutoencoderKLQwenImage) — Le variational auto-encoder de Qwen-Image (f8, 16 canaux latents) utilisé pour décoder les latents en images.
  • text_encoder (PreTrainedModel) — Un modèle Qwen3-VL (par exemple Qwen3VLModel de Qwen/Qwen3-VL-4B-Instruct). Le pipeline consomme une pile d’états cachés prélevés à partir de plusieurs couches décodeur, plutôt que l’état caché final.
  • tokenizer (AutoTokenizer) — Le tokenizer associé à l'encodeur de texte.
  • transformer (Krea2Transformer2DModel) — Le MMDiT à flux unique de Krea 2 qui prédit la vitesse de flow-matching.
  • text_encoder_select_layers (tuple[int, ...], facultatif) — Indices dans le tuple hidden_states de l’encodeur de texte (0 correspond à la sortie des embeddings) dont les états sont empilés par jeton pour servir de conditionnement textuel au transformeur. Doit contenir transformer.config.num_text_layers entrées.
  • is_distilled (bool, facultatif, par défaut False) — Indique si le transformer est un point de contrôle distillé en peu d'étapes (TDM/turbo). Si True, un décalage de pas de temps fixe mu=1.15 est utilisé ; sinon, mu est calculé à partir de la résolution de l'image.
  • patch_size (int, facultatif, valeur par défaut 2) — Côté des carrés de patches dans lesquels les latents sont regroupés avant d’entrer dans le transformeur. Le facteur de sous-échantillonnage effectif pixel-vers-token est vae_scale_factor * patch_size.

Le pipeline Krea 2 pour la génération d'images à partir de texte.

__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

Paramètres

  • prompt (str ou list[str], facultatif) — Le prompt ou les prompts pour guider la génération d’image. S’il n’est pas défini, il faut passer prompt_embeds.
  • negative_prompt (str ou list[str], facultatif) — Le ou les prompts à ne pas utiliser pour guider la génération d’image. Ignoré lorsque guidance_scale <= 0 ; par défaut, un prompt vide est utilisé lorsque le guidage est activé.
  • height (int, par défaut 1024) — La hauteur en pixels de l'image générée. Arrondie au multiple de 16 supérieur si nécessaire.
  • width (int, par défaut 1024) — La largeur en pixels de l'image générée. Arrondie au multiple de 16 supérieur si nécessaire.
  • num_inference_steps (int, par défaut 28) — Le nombre d'étapes de débruitage. Utilisez 28 pour le point de contrôle de base (midtrain) et 8 pour le point de contrôle distillé en peu d'étapes (TDM).
  • sigmas (list[float], facultatif) — Sigmas personnalisés pour le planificateur. S'ils ne sont pas définis, la grille par défaut linspace(1.0, 1/num_inference_steps, num_inference_steps) est utilisée (le décalage tenant compte de la résolution est appliqué à l'intérieur du planificateur).
  • guidance_scale (float, par défaut 4.5) — Échelle de guidage sans classifieur, selon la convention Krea 2 : la vitesse est calculée comme cond + guidance_scale * (cond - uncond) et le guidage est activé dès que guidance_scale > 0 (ce qui équivaut à la formulation habituelle du CFG avec une échelle de 1 + guidance_scale). Définir à 0.0 pour désactiver (par exemple pour le point de contrôle TDM).
  • num_images_per_prompt (int, par défaut 1) — Le nombre d'images à générer par invite.
  • generator (torch.Generator ou list[torch.Generator], facultatif) — Un ou plusieurs générateurs torch pour rendre la génération déterministe.
  • latents (torch.Tensor, facultatif) — Latents bruités pré-générés au format compact (batch_size, image_seq_len, in_channels), échantillonnés à partir d'une distribution gaussienne, à utiliser comme entrées pour la génération d'images.
  • prompt_embeds (torch.Tensor, facultatif) — Embeddings de texte pré-générés de forme (batch_size, text_seq_len, num_text_layers, text_hidden_dim). S'ils ne sont pas fournis, les embeddings sont générés à partir de prompt.
  • prompt_embeds_mask (torch.Tensor, facultatif) — Masque booléen pour prompt_embeds ; requis lorsque prompt_embeds est transmis.
  • negative_prompt_embeds (torch.Tensor, facultatif) — Embeddings de texte négatifs pré-générés ; même structure que prompt_embeds.
  • negative_prompt_embeds_mask (torch.Tensor, facultatif) — Masque booléen pour negative_prompt_embeds ; requis lorsque negative_prompt_embeds est transmis.
  • output_type (str, facultatif, par défaut "pil") — Le format de sortie de l’image générée. Choisissez entre "pil", "np", "pt" ou "latent".
  • return_dict (bool, facultatif, par défaut True) — Indique s'il faut retourner un Krea2PipelineOutput plutôt qu'un simple tuple.
  • callback_on_step_end (Callable, facultatif) — Une fonction appelée à la fin de chaque étape de débruitage avec callback_on_step_end(self, step, timestep, callback_kwargs).
  • callback_on_step_end_tensor_inputs (list[str], facultatif, par défaut ["latents"]) — La liste des entrées tensorielles pour la fonction callback_on_step_end. Doit être un sous-ensemble de ._callback_tensor_inputs.
  • attention_kwargs (dict, facultatif) — Un dictionnaire de kwargs qui, s'il est spécifié, est transmis à l'AttentionProcessor défini sous self.processor dans diffusers.models.attention_processor.
  • max_sequence_length (int, par défaut 512) — Longueur fixe de la séquence de texte consommée par le transformer ; les prompts sont complétés ou tronqués pour atteindre cette longueur.

Retourne

Krea2PipelineOutput ou tuple

Krea2PipelineOutput si return_dict est True, sinon un tuple dont le premier élément est une liste contenant les images générées.

Fonction appelée lors de l'appel du pipeline pour la génération.

Exemples :

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

Paramètres

  • prompt (str ou list[str], facultatif) — prompt à encoder
  • device — (torch.device) : torch device
  • num_images_per_prompt (int) — nombre d'images à générer par invite
  • prompt_embeds (torch.Tensor, facultatif) — Embeddings de texte pré-générés de forme (batch_size, text_seq_len, num_text_layers, text_hidden_dim). Peut être utilisé pour ajuster facilement les entrées de texte, par ex. le pondération des prompts. S'il n'est pas fourni, les embeddings de texte seront générés à partir de l'argument d'entrée prompt.
  • prompt_embeds_mask (torch.Tensor, facultatif) — Masque booléen pré-généré indiquant les jetons de texte valides, de forme (batch_size, text_seq_len). Requis lorsque prompt_embeds est transmis.
  • max_sequence_length (int, par défaut 512) — Longueur fixe de la séquence de texte consommée par le transformer ; les prompts sont complétés ou tronqués pour atteindre cette longueur.

get_text_hidden_states

< >

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

Tokeniser prompt dans la mise en page de longueur fixe Krea 2 et récupérer les états cachés de l'encodeur sélectionné.

Retourne un tuple (hidden_states, attention_mask) de formes (batch_size, text_seq_len, num_text_layers, text_hidden_dim) et (batch_size, text_seq_len) (bool).

prepare_position_ids

< >

( text_seq_len: intgrid_height: intgrid_width: intdevice: device )

Construisez les coordonnées rotationnelles (text_seq_len + grid_height * grid_width, 3) pour la séquence combinée : les jetons de texte sont placés à l’origine, tandis que les jetons d’image portent leurs coordonnées (0, h, w) de la grille latente.

Krea2PipelineOutput

classe diffusers.pipelines.krea2.Krea2PipelineOutput

< >

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

Paramètres

  • images (list[PIL.Image.Image] ou np.ndarray) — Liste d'images PIL débruitées de longueur batch_size ou tableau numpy de forme (batch_size, height, width, num_channels).

Classe de sortie pour le pipeline Krea 2.

Modulaire

Krea 2 est également disponible en tant que pipeline modulaire. Le guidage sans classifieur est configuré via le composant guider plutôt que par un argument d'appel guidance_scale. Krea 2 utilise un CFG ancré sur la condition, qui est ClassifierFreeGuidance avec use_original_formulation=True.

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

Nous fournissons également un exemple d’utilisation de Krea2 Turbo. Le point de contrôle distillé correspond à son propre ensemble de blocs (Krea2TurboAutoBlocks) : il fonctionne sans guidage (pas de guider), n’accepte pas de prompt négatif et échantillonne en quelques étapes. ModularPipeline.from_pretrained sélectionne automatiquement les blocs turbo à partir de la configuration is_distilled du point de contrôle, de sorte qu’aucune configuration de guidage n’est nécessaire :

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

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

Un ModularPipeline pour Krea 2.

Krea2AutoBlocks

classe diffusers.Krea2AutoBlocks

< >

( )

Pipeline Auto modulaire pour la génération d'images à partir de texte utilisant Krea 2 : encodage du texte -> débruitage principal (CFG symétrique) -> décodage.

Workflows pris en charge :

  • text2image : nécessite prompt

Composants : text_encoder (Qwen3VLModel) : l’encodeur de texte Qwen3-VL. tokenizer (AutoTokenizer) : le tokeniseur associé à l’encodeur de texte. guider (ClassifierFreeGuidance) transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Entrées : prompt (str) : Le prompt ou les prompts pour guider la génération d'images. negative_prompt (str, facultatif) : Le ou les prompts négatifs pour le CFG. max_sequence_length (int, facultatif, par défaut 512) : Longueur de séquence maximale pour l'encodage du prompt. num_images_per_prompt (int, facultatif, par défaut 1) : Le nombre d'images à générer par prompt. latents (Tensor, facultatif) : Latents bruités pré-générés pour la génération d'images. height (int, facultatif, par défaut 1024) : La hauteur en pixels de l'image générée. width (int, facultatif, par défaut 1024) : La largeur en pixels de l'image générée. generator (Generator, facultatif) : Générateur Torch pour une génération déterministe. num_inference_steps (int, facultatif, par défaut 28) : Le nombre d'étapes de débruitage. sigmas (list, facultatif) : Planification de sigmas personnalisée (par défaut, une rampe linéaire). attention_kwargs (dict, facultatif) : Arguments supplémentaires pour les processeurs d'attention. output_type (str, facultatif, par défaut pil) : Format de sortie : ‘pil’, ‘np’, ‘pt’.

Sorties : images (list) : Images générées.

Krea2TurboModularPipeline

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

Un ModularPipeline pour le point de contrôle distillé Krea 2 turbo (TDM). Il s'exécute sans guidage sans classifieur, il n'accepte donc pas de prompt négatif et ne possède aucun guideur.

Krea2TurboAutoBlocks

classe diffusers.Krea2TurboAutoBlocks

< >

( )

Pipeline Auto modulaire pour la génération d'images à partir de texte utilisant le point de contrôle distillé Krea 2 turbo : encodage du texte -> débruitage principal (sans guidage) -> décodage.

Workflows pris en charge :

  • text2image : nécessite prompt

Composants : text_encoder (Qwen3VLModel) : l’encodeur de texte Qwen3-VL. tokenizer (AutoTokenizer) : le tokeniseur associé à l’encodeur de texte. transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Entrées : prompt (str) : Le prompt ou les prompts pour guider la génération d'images. max_sequence_length (int, facultatif, par défaut 512) : Longueur de séquence maximale pour l'encodage du prompt. num_images_per_prompt (int, facultatif, par défaut 1) : Le nombre d'images à générer par prompt. latents (Tensor, facultatif) : Latents bruités pré-générés pour la génération d'images. height (int, facultatif, par défaut 1024) : La hauteur en pixels de l'image générée. width (int, facultatif, par défaut 1024) : La largeur en pixels de l'image générée. generator (Generator, facultatif) : Générateur Torch pour une génération déterministe. num_inference_steps (int, facultatif, par défaut 8) : Le nombre d'étapes de débruitage. sigmas (list, facultatif) : Planification de sigma personnalisée (par défaut, une rampe linéaire). attention_kwargs (dict, facultatif) : Arguments supplémentaires pour les processeurs d'attention. output_type (str, facultatif, par défaut pil) : Format de sortie : ‘pil’, ‘np’, ‘pt’.

Sorties : images (list) : Images générées.

Update on GitHub