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