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) é um modelo de texto para imagem baseado em flow-matching, construído em torno de um MMDiT de fluxo único com grouped-query attention. Um codificador de texto Qwen3-VL fornece o condicionamento: em vez do último estado oculto, estados ocultos de doze camadas decodificadoras são extraídos por token e fundidos dentro do transformer por um pequeno estágio de fusão de texto. As imagens são decodificadas com o VAE do Qwen-Image.

Dois checkpoints são liberados, compartilhando a mesma arquitetura, mas com configurações de sampler recomendadas diferentes:

  • Base (midtrain) — use o sampler completo com classifier-free guidance: num_inference_steps=28, guidance_scale=4.5.
  • TDM (destilado) — destilado para amostragem em poucas etapas, execute com num_inference_steps=8 e orientação desativada (guidance_scale=0.0).

guidance_scale segue a convenção do Krea 2: a velocidade é calculada como cond + guidance_scale * (cond - uncond) e o guidance é ativado sempre que guidance_scale > 0 (isso equivale à formulação usual de CFG com escala 1 + guidance_scale).

Texto para imagem

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

Fornecemos, adicionalmente, um exemplo para o uso do 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")

Carregando checkpoints de arquivo único

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 )

Parâmetros

  • scheduler (FlowMatchEulerDiscreteScheduler) — Agendador de fluxo-Euler. O agendamento de sigma do Krea 2 é o deslocamento exponencial de tempo consciente da resolução, portanto a configuração do agendador deve definir use_dynamic_shifting=True junto com os parâmetros de deslocamento do Krea 2 (base_shift=0.5, max_shift=1.15, base_image_seq_len=256, max_image_seq_len=6400).
  • vae (AutoencoderKLQwenImage) — O autoencoder variacional do Qwen-Image (f8, 16 canais latentes) usado para decodificar latentes em imagens.
  • text_encoder (PreTrainedModel) — Um modelo Qwen3-VL (por exemplo, Qwen3VLModel de Qwen/Qwen3-VL-4B-Instruct). O pipeline consome uma pilha de hidden states obtidos de várias camadas decodificadoras, em vez do último hidden state.
  • tokenizer (AutoTokenizer) — O tokenizer associado ao codificador de texto.
  • transformer (Krea2Transformer2DModel) — O MMDiT de fluxo único do Krea 2 que prevê a velocidade do flow-matching.
  • text_encoder_select_layers (tuple[int, ...], opcional) — Índices da tupla hidden_states do codificador de texto (0 é a saída do embedding) cujos estados são empilhados por token como a condição de texto do transformer. Deve ter transformer.config.num_text_layers entradas.
  • is_distilled (bool, opcional, padrão False) — Indica se o transformer é o checkpoint destilado de poucas etapas (TDM/turbo). Quando True, um deslocamento de passo de tempo fixo mu=1.15 é usado; caso contrário, mu é calculado a partir da resolução da imagem.
  • patch_size (int, opcional, padrão 2) — Comprimento do lado dos patches quadrados nos quais os latentes são empacotados antes de entrarem no transformer. O fator efetivo de downsampling de pixel para token é vae_scale_factor * patch_size.

O pipeline Krea 2 para geração de imagens a partir de texto.

__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

Parâmetros

  • prompt (str ou list[str], opcional) — O prompt ou os prompts para orientar a geração da imagem. Se não for definido, é necessário passar prompt_embeds.
  • negative_prompt (str ou list[str], opcional) — O prompt ou os prompts que não devem orientar a geração da imagem. Ignorado quando guidance_scale <= 0; padrão para um prompt vazio quando a orientação está habilitada.
  • height (int, padrão 1024) — A altura em pixels da imagem gerada. Arredondada para cima para um múltiplo de 16, se necessário.
  • width (int, padrão 1024) — A largura em pixels da imagem gerada. Arredondada para cima para um múltiplo de 16, se necessário.
  • num_inference_steps (int, padrão 28) — O número de etapas de remoção de ruído. Use 28 para o checkpoint base (midtrain) e 8 para o checkpoint destilado de poucas etapas (TDM).
  • sigmas (list[float], opcional) — Sigmas personalizados para o agendador. Se não definidos, a grade padrão linspace(1.0, 1/num_inference_steps, num_inference_steps) é utilizada (o deslocamento consciente da resolução é aplicado dentro do agendador).
  • guidance_scale (float, padrão 4.5) — Escala de orientação sem classificador, seguindo a convenção do Krea 2: a velocidade é calculada como cond + guidance_scale * (cond - uncond) e a orientação é ativada sempre que guidance_scale > 0 (isso equivale à formulação usual de CFG com escala 1 + guidance_scale). Defina como 0.0 para desativar (por exemplo, para o checkpoint TDM).
  • num_images_per_prompt (int, padrão 1) — O número de imagens a serem geradas por prompt.
  • generator (torch.Generator ou list[torch.Generator], opcional) — Um ou mais geradores do torch para tornar a geração determinística.
  • latents (torch.Tensor, opcional) — Latentes ruidosos pré-gerados em formato empacotado (batch_size, image_seq_len, in_channels), amostrados de uma distribuição gaussiana, a serem usados como entradas para a geração de imagens.
  • prompt_embeds (torch.Tensor, opcional) — Incorporações de texto pré-geradas com forma (batch_size, text_seq_len, num_text_layers, text_hidden_dim). Se não fornecidas, as incorporações são geradas a partir de prompt.
  • prompt_embeds_mask (torch.Tensor, opcional) — Máscara booleana para prompt_embeds; obrigatória quando prompt_embeds é fornecida.
  • negative_prompt_embeds (torch.Tensor, opcional) — Incorporações de texto negativo pré-geradas; mesmo layout que prompt_embeds.
  • negative_prompt_embeds_mask (torch.Tensor, opcional) — Máscara booleana para negative_prompt_embeds; obrigatória quando negative_prompt_embeds é fornecida.
  • output_type (str, opcional, padrão "pil") — O formato de saída da imagem gerada. Escolha entre "pil", "np", "pt" ou "latent".
  • return_dict (bool, opcional, padrão True) — Se deve ou não retornar um Krea2PipelineOutput em vez de uma tupla simples.
  • callback_on_step_end (Callable, opcional) — Uma função que é chamada no final de cada etapa de remoção de ruído com callback_on_step_end(self, step, timestep, callback_kwargs).
  • callback_on_step_end_tensor_inputs (list[str], opcional, padrão ["latents"]) — A lista de entradas de tensor para a função callback_on_step_end. Deve ser um subconjunto de ._callback_tensor_inputs.
  • attention_kwargs (dict, opcional) — Um dicionário de kwargs que, se especificado, é passado ao AttentionProcessor conforme definido em self.processor em diffusers.models.attention_processor.
  • max_sequence_length (int, padrão 512) — Comprimento fixo da sequência de texto consumido pelo transformer; os prompts são preenchidos ou truncados até esse valor.

Retorna

Krea2PipelineOutput ou tuple

Krea2PipelineOutput se return_dict for True, caso contrário uma tuple, cujo primeiro elemento é uma lista com as imagens geradas.

Função invocada ao chamar o pipeline para geração.

Exemplos:

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

Parâmetros

  • prompt (str ou list[str], opcional) — prompt a ser codificado
  • device — (torch.device): dispositivo torch
  • num_images_per_prompt (int) — número de imagens que devem ser geradas por prompt
  • prompt_embeds (torch.Tensor, opcional) — Incorporações de texto pré-geradas com a forma (batch_size, text_seq_len, num_text_layers, text_hidden_dim). Pode ser usado para ajustar facilmente as entradas de texto, por exemplo, ponderação de prompt. Se não for fornecido, as incorporações de texto serão geradas a partir do argumento de entrada prompt.
  • prompt_embeds_mask (torch.Tensor, opcional) — Máscara booleana pré-gerada que marca os tokens de texto válidos, com forma (batch_size, text_seq_len). Obrigatória quando prompt_embeds é fornecida.
  • max_sequence_length (int, padrão 512) — Comprimento fixo da sequência de texto consumido pelo transformer; os prompts são preenchidos ou truncados até esse valor.

get_text_hidden_states

< >

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

Tokenize prompt no layout de comprimento fixo do Krea 2 e capture os estados ocultos do encoder selecionado.

Retorna uma tupla (hidden_states, attention_mask) de formas (batch_size, text_seq_len, num_text_layers, text_hidden_dim) e (batch_size, text_seq_len) (bool).

prepare_position_ids

< >

( text_seq_len: intgrid_height: intgrid_width: intdevice: device )

Construa as coordenadas rotacionais (text_seq_len + grid_height * grid_width, 3) para a sequência combinada: os tokens de texto ficam na origem, e os tokens de imagem carregam suas coordenadas (0, h, w) da grade latente.

Krea2PipelineOutput

classe diffusers.pipelines.krea2.Krea2PipelineOutput

< >

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

Parâmetros

  • images (list[PIL.Image.Image] ou np.ndarray) — Lista de imagens PIL desruidas de comprimento batch_size ou array numpy de forma (batch_size, height, width, num_channels).

Classe de saída para o pipeline Krea 2.

Modular

Krea 2 também está disponível como um pipeline modular. A classifier-free guidance é configurada por meio do componente guider em vez de um argumento de chamada guidance_scale. O Krea 2 usa CFG ancorado em cond, que é ClassifierFreeGuidance com 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")

Também fornecemos um exemplo para o uso do Krea2 Turbo. O checkpoint destilado mapeia para seu próprio conjunto de blocos (Krea2TurboAutoBlocks): ele executa sem orientação (sem guider), não aceita prompt negativo e amostra em poucas etapas. ModularPipeline.from_pretrained seleciona automaticamente os blocos turbo a partir da configuração is_distilled do checkpoint, portanto nenhuma configuração de orientação é necessária:

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 )

Uma ModularPipeline para o Krea 2.

Krea2AutoBlocks

classe diffusers.Krea2AutoBlocks

< >

( )

Pipeline modular automático para geração de texto para imagem usando Krea 2: codificar texto -> desruído principal (CFG simétrico) -> decodificar.

Fluxos de trabalho suportados:

  • text2image: requer prompt

Componentes: text_encoder (Qwen3VLModel): O codificador de texto do Qwen3-VL. tokenizer (AutoTokenizer): O tokenizador emparelhado com o codificador de texto. guider (ClassifierFreeGuidance) transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Entradas: prompt (str): O prompt ou prompts para orientar a geração de imagem. negative_prompt (str, opcional): O(s) prompt(s) negativo(s) para CFG. max_sequence_length (int, opcional, padrão 512): Comprimento máximo da sequência para a codificação do prompt. num_images_per_prompt (int, opcional, padrão 1): O número de imagens a serem geradas por prompt. latents (Tensor, opcional): Latentes ruidosos pré-gerados para a geração de imagem. height (int, opcional, padrão 1024): A altura em pixels da imagem gerada. width (int, opcional, padrão 1024): A largura em pixels da imagem gerada. generator (Generator, opcional): Gerador do Torch para geração determinística. num_inference_steps (int, opcional, padrão 28): O número de etapas de remoção de ruído. sigmas (list, opcional): Agenda de sigmas personalizada (padrão: rampa linear). attention_kwargs (dict, opcional): Kwargs adicionais para os processadores de atenção. output_type (str, opcional, padrão pil): Formato de saída: ‘pil’, ‘np’, ‘pt’.

Saídas: imagens (list): Imagens geradas.

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 )

Uma ModularPipeline para o checkpoint destilado Krea 2 turbo (TDM). Ele é executado sem classifier-free guidance, portanto não aceita prompt negativo e não possui guider.

Krea2TurboAutoBlocks

classe diffusers.Krea2TurboAutoBlocks

< >

( )

Pipeline Auto Modular para geração de texto para imagem usando o checkpoint Krea 2 turbo destilado: codificar texto -> desruidar núcleo (sem orientação) -> decodificar.

Fluxos de trabalho suportados:

  • text2image: requer prompt

Componentes: text_encoder (Qwen3VLModel): O codificador de texto do Qwen3-VL. tokenizer (AutoTokenizer): O tokenizador emparelhado com o codificador de texto. transformer (Krea2Transformer2DModel) scheduler (FlowMatchEulerDiscreteScheduler) vae (AutoencoderKLQwenImage) image_processor (VaeImageProcessor)

Entradas: prompt (str): O prompt ou prompts para orientar a geração de imagens. max_sequence_length (int, opcional, padrão 512): Comprimento máximo da sequência para a codificação do prompt. num_images_per_prompt (int, opcional, padrão 1): O número de imagens a serem geradas por prompt. latents (Tensor, opcional): Latentes ruidosos pré-gerados para a geração de imagens. height (int, opcional, padrão 1024): A altura em pixels da imagem gerada. width (int, opcional, padrão 1024): A largura em pixels da imagem gerada. generator (Generator, opcional): Gerador do Torch para geração determinística. num_inference_steps (int, opcional, padrão 8): O número de etapas de remoção de ruído. sigmas (list, opcional): Agenda de sigmas personalizada (padrão: rampa linear). attention_kwargs (dict, opcional): Kwargs adicionais para os processadores de atenção. output_type (str, opcional, padrão pil): Formato de saída: ‘pil’, ‘np’, ‘pt’.

Saídas: imagens (list): Imagens geradas.

Update on GitHub