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) 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 = ['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