From 0b05beba688352f160362ebf57dd2baf589798cc Mon Sep 17 00:00:00 2001 From: lucasruan1618 Date: Mon, 3 Aug 2026 15:43:01 +0000 Subject: [PATCH 1/2] [Modular]: Add Krea2 Image2Image, Inpaint, References modular pipelines --- docs/source/en/api/pipelines/krea2.md | 103 +++ .../models/transformers/transformer_krea2.py | 67 +- .../modular_pipelines/krea2/before_denoise.py | 488 +++++++++++ .../modular_pipelines/krea2/decoders.py | 110 ++- .../modular_pipelines/krea2/denoise.py | 344 ++++++++ .../modular_pipelines/krea2/encoders.py | 716 +++++++++++++++- .../krea2/modular_blocks_krea2.py | 762 +++++++++++++++++- .../krea2/modular_blocks_krea2_turbo.py | 481 ++++++++++- .../test_models_transformer_krea2.py | 31 + .../krea2/test_modular_pipeline_krea2.py | 79 +- .../test_modular_pipeline_krea2_turbo.py | 79 +- 11 files changed, 3188 insertions(+), 72 deletions(-) diff --git a/docs/source/en/api/pipelines/krea2.md b/docs/source/en/api/pipelines/krea2.md index b6e6fd5998c3..2142cdc88124 100644 --- a/docs/source/en/api/pipelines/krea2.md +++ b/docs/source/en/api/pipelines/krea2.md @@ -106,6 +106,109 @@ image = pipe( image.save("krea2.png") ``` +The same modular pipeline automatically selects image-to-image generation when `image` is provided. `strength` +controls how strongly the result can depart from the source image. + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png") +image = pipe( + prompt="a cat wearing a knitted wizard hat", + image=init_image, + height=init_image.height, + width=init_image.width, + strength=0.8, + num_inference_steps=28, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("krea2_img2img.png") +``` + +Provide both `image` and `mask_image` to select inpainting. White mask pixels are regenerated and black mask pixels +are preserved. + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("krea/Krea-2-Raw") +pipe.load_components(dtype=torch.bfloat16) +pipe.to("cuda") + +init_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint.png") +mask_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/inpaint_mask.png") +image = pipe( + prompt="a small red fox sitting on a park bench", + image=init_image, + mask_image=mask_image, + height=init_image.height, + width=init_image.width, + strength=0.9, + num_inference_steps=28, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("krea2_inpaint.png") +``` + +### Reference-conditioned generation + +Pass `reference_image` to condition generation on clean reference-image tokens and an image-grounded Qwen3-VL prompt +encoding. Unlike conventional image-to-image generation, the target starts from pure noise, so this workflow does not +use `strength`. It is intended for LoRAs trained with the same reference-conditioning layout and is not tied to one +specific identity or editing adapter. + +The following example uses the community [Krea 2 Identity Edit](https://huggingface.co/conradlocke/krea2-identity-edit) +LoRA: + +```python +import torch +from diffusers import ModularPipeline +from diffusers.utils import load_image + +pipe = ModularPipeline.from_pretrained("krea/Krea-2-Turbo") +pipe.load_components(dtype=torch.bfloat16) +pipe.load_lora_weights( + "conradlocke/krea2-identity-edit", + weight_name="krea2_identity_edit_v1_2_r64.safetensors", + adapter_name="krea2_edit", +) +pipe.to("cuda") + +scene_image = load_image( + "https://raw.githubusercontent.com/lucasruan1618/Image_storage/main/Input/cute_dog.png" +) +subject_image = load_image( + "https://raw.githubusercontent.com/lucasruan1618/Image_storage/main/Input/cute_cat.png" +) +image = pipe( + prompt="place the wizard cat from the second image sitting on the bench beside the dog from the first image", + reference_image=scene_image, + reference_image_2=subject_image, + height=1024, + width=1024, + reference_image_encoder_resolution=768, + reference_attention_scale=[1.0, 4.0], + num_inference_steps=10, + generator=torch.Generator("cuda").manual_seed(0), +).images[0] +image.save("krea2_reference.png") +``` + +For two-reference generation, `reference_image` is the scene and `reference_image_2` is the subject, matching the +adapter's training order. `reference_image_encoder_resolution` controls the maximum reference-image side length passed +to Qwen3-VL. `reference_attention_scale` accepts either one value for all references or one value per reference; the +example leaves scene attention unchanged and boosts subject fidelity. The adapter's recommended LoRA scale is `1.0`. +References are resized to the requested output dimensions before VAE encoding, so use similar aspect ratios to avoid +distortion. + We additionally provide an example for using Krea2 Turbo. The distilled checkpoint maps to its own set of blocks ([`Krea2TurboAutoBlocks`]): it runs guidance-free (no `guider`), takes no negative prompt, and samples in a few steps. `ModularPipeline.from_pretrained` picks the turbo blocks automatically from the checkpoint's `is_distilled` config, so diff --git a/src/diffusers/models/transformers/transformer_krea2.py b/src/diffusers/models/transformers/transformer_krea2.py index d1f6cd0ecded..f96464dd3cdc 100644 --- a/src/diffusers/models/transformers/transformer_krea2.py +++ b/src/diffusers/models/transformers/transformer_krea2.py @@ -74,12 +74,19 @@ def __call__( query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1) key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1) + enable_gqa = attn.num_heads != attn.num_kv_heads + if attention_mask is not None and attention_mask.dtype != torch.bool and enable_gqa: + repeats = attn.num_heads // attn.num_kv_heads + key = key.repeat_interleave(repeats, dim=2) + value = value.repeat_interleave(repeats, dim=2) + enable_gqa = False + hidden_states = dispatch_attention_fn( query, key, value, attn_mask=attention_mask, - enable_gqa=attn.num_heads != attn.num_kv_heads, + enable_gqa=enable_gqa, backend=self._attention_backend, parallel_config=self._parallel_config, ) @@ -452,6 +459,8 @@ def forward( timestep: torch.Tensor, position_ids: torch.Tensor, encoder_attention_mask: torch.Tensor | None = None, + reference_hidden_states: list[torch.Tensor] | None = None, + reference_attention_scale: float | list[float] = 1.0, attention_kwargs: dict[str, Any] | None = None, return_dict: bool = True, ) -> Transformer2DModelOutput | tuple[torch.Tensor]: @@ -470,6 +479,12 @@ def forward( latent-grid coordinates. encoder_attention_mask (`torch.Tensor` of shape `(batch_size, text_seq_len)`, *optional*): Boolean mask marking valid text tokens. Pass `None` when every text token is valid. + reference_hidden_states (`list[torch.Tensor]`, *optional*): + Packed clean reference-image latents prepended in list order before the noisy image tokens. Each tensor + has shape `(batch_size, reference_seq_len, in_channels)`. + reference_attention_scale (`float` or `list[float]`, *optional*, defaults to `1.0`): + Multiplier applied to target-token attention probabilities for each reference-image block. A float is + applied to every reference; a list sets one multiplier per reference in the same order. attention_kwargs (`dict`, *optional*): A kwargs dictionary that, when it contains a `scale` entry, sets the LoRA scale applied to this transformer's adapters for the duration of the forward pass. @@ -485,6 +500,32 @@ def forward( batch_size, image_seq_len, _ = hidden_states.shape text_seq_len = encoder_hidden_states.shape[1] + reference_seq_lens = [] if reference_hidden_states is None else [x.shape[1] for x in reference_hidden_states] + reference_seq_len = sum(reference_seq_lens) + + if isinstance(reference_attention_scale, list): + reference_attention_scales = reference_attention_scale + elif reference_hidden_states is None: + reference_attention_scales = [] + else: + reference_attention_scales = [reference_attention_scale] * len(reference_hidden_states) + if reference_hidden_states is None and reference_attention_scale != 1.0: + raise ValueError("`reference_attention_scale` requires `reference_hidden_states`.") + if reference_hidden_states is not None and len(reference_hidden_states) == 0: + raise ValueError("`reference_hidden_states` must contain at least one tensor.") + if len(reference_attention_scales) != len(reference_seq_lens): + raise ValueError( + "`reference_attention_scale` must contain one value per reference tensor, but got " + f"{len(reference_attention_scales)} values for {len(reference_seq_lens)} references." + ) + if any(scale < 0 for scale in reference_attention_scales): + raise ValueError(f"`reference_attention_scale` must be non-negative, but is {reference_attention_scale}.") + sequence_length = text_seq_len + reference_seq_len + image_seq_len + if position_ids.shape[0] != sequence_length: + raise ValueError( + f"`position_ids` has sequence length {position_ids.shape[0]}, but the combined text, reference, and " + f"image sequence has length {sequence_length}." + ) temb = self.time_embed(timestep, dtype=hidden_states.dtype) temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh")) @@ -495,14 +536,32 @@ def forward( # Key-padding masks of shape (B, 1, 1, L): padded text tokens are excluded as attention keys everywhere; # their own (garbage) lanes are never read back and are dropped at the output slice. text_attention_mask = encoder_attention_mask[:, None, None, :] - image_mask = encoder_attention_mask.new_ones((batch_size, image_seq_len)) + image_mask = encoder_attention_mask.new_ones((batch_size, reference_seq_len + image_seq_len)) attention_mask = torch.cat([encoder_attention_mask, image_mask], dim=1)[:, None, None, :] encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask) encoder_hidden_states = self.txt_in(encoder_hidden_states) hidden_states = self.img_in(hidden_states) - hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + if reference_hidden_states is not None: + reference_hidden_states = [self.img_in(x) for x in reference_hidden_states] + hidden_states = torch.cat([encoder_hidden_states, *reference_hidden_states, hidden_states], dim=1) + else: + hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + if reference_hidden_states is not None and any(scale != 1.0 for scale in reference_attention_scales): + reference_attention_bias = hidden_states.new_zeros((batch_size, 1, sequence_length, sequence_length)) + target_start = text_seq_len + reference_seq_len + reference_start = text_seq_len + for reference_length, scale in zip(reference_seq_lens, reference_attention_scales): + reference_end = reference_start + reference_length + reference_attention_bias[:, :, target_start:, reference_start:reference_end] = math.log( + max(scale, 1e-4) + ) + reference_start = reference_end + if attention_mask is not None: + reference_attention_bias.masked_fill_(~attention_mask, float("-inf")) + attention_mask = reference_attention_bias image_rotary_emb = self.rotary_emb(position_ids) @@ -514,7 +573,7 @@ def forward( else: hidden_states = block(hidden_states, temb_mod, image_rotary_emb, attention_mask) - hidden_states = hidden_states[:, text_seq_len:] + hidden_states = hidden_states[:, -image_seq_len:] output = self.final_layer(hidden_states, temb) if not return_dict: diff --git a/src/diffusers/modular_pipelines/krea2/before_denoise.py b/src/diffusers/modular_pipelines/krea2/before_denoise.py index 63810d30a903..62600ea9a9dc 100644 --- a/src/diffusers/modular_pipelines/krea2/before_denoise.py +++ b/src/diffusers/modular_pipelines/krea2/before_denoise.py @@ -251,6 +251,203 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi return components, state +# auto_docstring +class Krea2ImageInputsStep(ModularPipelineBlocks): + """ + Pack image latents into Krea 2 image tokens and expand image and mask inputs to the effective prompt batch. + + Inputs: + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + Effective batch size. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + processed_mask_image (`Tensor`): + The batch-expanded inpainting mask. + height (`int`): + The generation height inferred from the image. + width (`int`): + The generation width inferred from the image. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Pack image latents into Krea 2 image tokens and expand image and mask inputs to the effective prompt " + "batch." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image_latents"), + InputParam( + name="processed_mask_image", type_hint=torch.Tensor, description="The preprocessed inpainting mask." + ), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("num_images_per_prompt", default=1), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("image_latents"), + OutputParam( + name="processed_mask_image", type_hint=torch.Tensor, description="The batch-expanded inpainting mask." + ), + OutputParam(name="height", type_hint=int, description="The generation height inferred from the image."), + OutputParam(name="width", type_hint=int, description="The generation width inferred from the image."), + ] + + @staticmethod + def repeat_to_batch_size(input_name, input_tensor, prompt_batch_size, num_images_per_prompt): + if input_tensor.shape[0] == 1: + repeat_by = prompt_batch_size * num_images_per_prompt + elif input_tensor.shape[0] == prompt_batch_size: + repeat_by = num_images_per_prompt + else: + raise ValueError( + f"`{input_name}` must have batch size 1 or {prompt_batch_size}, but got {input_tensor.shape[0]}" + ) + return input_tensor.repeat_interleave(repeat_by, dim=0) + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + image_latents = block_state.image_latents + if image_latents.ndim != 5 or image_latents.shape[2] != 1: + raise ValueError( + f"`image_latents` must have shape (batch, channels, 1, height, width), got {image_latents.shape}" + ) + + image_height = image_latents.shape[-2] * components.vae_scale_factor + image_width = image_latents.shape[-1] * components.vae_scale_factor + block_state.height = block_state.height or image_height + block_state.width = block_state.width or image_width + if block_state.height != image_height or block_state.width != image_width: + raise ValueError( + f"The encoded image is {image_height}x{image_width}, but the requested output is " + f"{block_state.height}x{block_state.width}." + ) + + p = components.patch_size + batch_size, channels, _, latent_height, latent_width = image_latents.shape + image_latents = image_latents[:, :, 0].view(batch_size, channels, latent_height // p, p, latent_width // p, p) + image_latents = image_latents.permute(0, 2, 4, 1, 3, 5).reshape( + batch_size, (latent_height // p) * (latent_width // p), channels * p * p + ) + + prompt_batch_size = block_state.batch_size // block_state.num_images_per_prompt + block_state.image_latents = self.repeat_to_batch_size( + "image_latents", image_latents, prompt_batch_size, block_state.num_images_per_prompt + ) + if block_state.processed_mask_image is not None: + block_state.processed_mask_image = self.repeat_to_batch_size( + "processed_mask_image", + block_state.processed_mask_image, + prompt_batch_size, + block_state.num_images_per_prompt, + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2ReferenceInputsStep(ModularPipelineBlocks): + """ + Pack reference-image latents and expand them to the effective prompt batch. + + Inputs: + reference_image_latents (`list`): + Normalized reference-image latents from the VAE encoder in conditioning order. + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + Effective batch size. + + Outputs: + reference_image_latents (`list`): + Packed reference-image latents expanded to the effective batch in conditioning order. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Pack reference-image latents and expand them to the effective prompt batch." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="reference_image_latents", + required=True, + type_hint=list[torch.Tensor], + description="Normalized reference-image latents from the VAE encoder in conditioning order.", + ), + InputParam.template("num_images_per_prompt", default=1), + InputParam(name="batch_size", required=True, type_hint=int, description="Effective batch size."), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="reference_image_latents", + type_hint=list[torch.Tensor], + description="Packed reference-image latents expanded to the effective batch in conditioning order.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + p = components.patch_size + prompt_batch_size = block_state.batch_size // block_state.num_images_per_prompt + packed_reference_image_latents = [] + for reference_image_latents in block_state.reference_image_latents: + if reference_image_latents.ndim != 5 or reference_image_latents.shape[2] != 1: + raise ValueError( + "Each `reference_image_latents` tensor must have shape (batch, channels, 1, height, width), but " + f"got {reference_image_latents.shape}." + ) + batch_size, channels, _, latent_height, latent_width = reference_image_latents.shape + reference_image_latents = reference_image_latents[:, :, 0].view( + batch_size, channels, latent_height // p, p, latent_width // p, p + ) + reference_image_latents = reference_image_latents.permute(0, 2, 4, 1, 3, 5).reshape( + batch_size, (latent_height // p) * (latent_width // p), channels * p * p + ) + if batch_size == 1: + repeat_by = prompt_batch_size * block_state.num_images_per_prompt + elif batch_size == prompt_batch_size: + repeat_by = block_state.num_images_per_prompt + else: + raise ValueError( + f"Each reference must have batch size 1 or {prompt_batch_size}, but got {batch_size}." + ) + packed_reference_image_latents.append(reference_image_latents.repeat_interleave(repeat_by, dim=0)) + block_state.reference_image_latents = packed_reference_image_latents + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Krea2PrepareLatentsStep(ModularPipelineBlocks): """ @@ -360,6 +557,211 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi return components, state +# auto_docstring +class Krea2ApplyStrengthStep(ModularPipelineBlocks): + """ + Truncate the Krea 2 denoising schedule according to image-to-image or inpainting strength. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + num_inference_steps (`int`): + The number of denoising steps. + timesteps (`Tensor`): + The full denoising schedule. + + Outputs: + timesteps (`Tensor`): + The strength-adjusted timesteps. + num_inference_steps (`int`): + The strength-adjusted denoising step count. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Truncate the Krea 2 denoising schedule according to image-to-image or inpainting strength." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("strength", default=0.9), + InputParam.template("num_inference_steps", required=True), + InputParam( + name="timesteps", required=True, type_hint=torch.Tensor, description="The full denoising schedule." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="timesteps", type_hint=torch.Tensor, description="The strength-adjusted timesteps."), + OutputParam( + name="num_inference_steps", type_hint=int, description="The strength-adjusted denoising step count." + ), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + if block_state.strength < 0 or block_state.strength > 1: + raise ValueError(f"`strength` must be in [0.0, 1.0], but is {block_state.strength}") + init_timestep = min(block_state.num_inference_steps * block_state.strength, block_state.num_inference_steps) + t_start = int(max(block_state.num_inference_steps - init_timestep, 0)) + begin_index = t_start * components.scheduler.order + block_state.timesteps = block_state.timesteps[begin_index:] + block_state.num_inference_steps -= t_start + if block_state.num_inference_steps < 1: + raise ValueError( + f"After applying `strength={block_state.strength}`, the number of denoising steps is " + f"{block_state.num_inference_steps}, but it must be at least 1." + ) + components.scheduler.set_begin_index(begin_index) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2PrepareImageLatentsStep(ModularPipelineBlocks): + """ + Add noise at the first selected timestep to packed Krea 2 image latents. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + latents (`Tensor`): + Pre-generated noisy latents for image generation. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + timesteps (`Tensor`): + The selected denoising timesteps. + + Outputs: + initial_noise (`Tensor`): + The sampled initial noise. + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Add noise at the first selected timestep to packed Krea 2 image latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam.template("image_latents", required=True), + InputParam( + name="timesteps", + required=True, + type_hint=torch.Tensor, + description="The selected denoising timesteps.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="initial_noise", type_hint=torch.Tensor, description="The sampled initial noise."), + OutputParam.template("latents"), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + if block_state.image_latents.shape != block_state.latents.shape: + raise ValueError( + f"`image_latents` and `latents` must have the same shape, got " + f"{block_state.image_latents.shape} and {block_state.latents.shape}" + ) + latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0]) + block_state.initial_noise = block_state.latents + block_state.latents = components.scheduler.scale_noise( + block_state.image_latents, latent_timestep, block_state.initial_noise + ) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2PrepareMaskLatentsStep(ModularPipelineBlocks): + """ + Resize and pack a preprocessed inpainting mask into Krea 2 image-token space. + + Inputs: + processed_mask_image (`Tensor`): + The preprocessed inpainting mask. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + mask (`Tensor`): + The packed latent-space mask. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Resize and pack a preprocessed inpainting mask into Krea 2 image-token space." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_mask_image", + required=True, + type_hint=torch.Tensor, + description="The preprocessed inpainting mask.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam.template("dtype"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam(name="mask", type_hint=torch.Tensor, description="The packed latent-space mask.")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + p = components.patch_size + latent_height = block_state.height // components.vae_scale_factor + latent_width = block_state.width // components.vae_scale_factor + mask = torch.nn.functional.interpolate( + block_state.processed_mask_image, size=(latent_height, latent_width), mode="nearest" + ) + channels = components.transformer.config.in_channels // (p**2) + mask = mask.repeat(1, channels, 1, 1).to(device=components._execution_device, dtype=block_state.dtype) + batch_size = mask.shape[0] + mask = mask.view(batch_size, channels, latent_height // p, p, latent_width // p, p) + mask = mask.permute(0, 2, 4, 1, 3, 5) + block_state.mask = mask.reshape(batch_size, (latent_height // p) * (latent_width // p), channels * p * p) + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Krea2SetTimestepsStep(ModularPipelineBlocks): """ @@ -588,3 +990,89 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi self.set_block_state(state, block_state) return components, state + + +# auto_docstring +class Krea2PrepareReferencePositionIdsStep(ModularPipelineBlocks): + """ + Build rotary position ids for a [text | reference | target] Krea 2 sequence. + + Inputs: + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + prompt_embeds (`Tensor`): + Batch-expanded text features. + reference_image_latents (`list`): + Packed reference-image latents in conditioning order. + + Outputs: + position_ids (`Tensor`): + Rotary coordinates for the [text | reference | target] sequence. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Build rotary position ids for a [text | reference | target] Krea 2 sequence." + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam( + name="prompt_embeds", + required=True, + type_hint=torch.Tensor, + description="Batch-expanded text features.", + ), + InputParam( + name="reference_image_latents", + required=True, + type_hint=list[torch.Tensor], + description="Packed reference-image latents in conditioning order.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="position_ids", + type_hint=torch.Tensor, + description="Rotary coordinates for the [text | reference | target] sequence.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + p = components.patch_size + grid_height = block_state.height // (components.vae_scale_factor * p) + grid_width = block_state.width // (components.vae_scale_factor * p) + image_seq_len = grid_height * grid_width + if any(reference.shape[1] != image_seq_len for reference in block_state.reference_image_latents): + reference_lengths = [reference.shape[1] for reference in block_state.reference_image_latents] + raise ValueError( + "Each packed reference image and the target must have the same token count, but got reference " + f"lengths {reference_lengths} and target length {image_seq_len}." + ) + + text_ids = torch.zeros(block_state.prompt_embeds.shape[1], 3, device=device) + image_ids = torch.zeros(grid_height, grid_width, 3, device=device) + image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None] + image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :] + reference_ids = [] + for frame in range(1, len(block_state.reference_image_latents) + 1): + ids = image_ids.clone() + ids[..., 0] = frame + reference_ids.append(ids.reshape(image_seq_len, 3)) + target_ids = image_ids.clone() + target_ids[..., 0] = 0 + block_state.position_ids = torch.cat([text_ids, *reference_ids, target_ids.reshape(image_seq_len, 3)], dim=0) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/krea2/decoders.py b/src/diffusers/modular_pipelines/krea2/decoders.py index fd308b5ef648..4bde70475c73 100644 --- a/src/diffusers/modular_pipelines/krea2/decoders.py +++ b/src/diffusers/modular_pipelines/krea2/decoders.py @@ -16,7 +16,7 @@ import torch from ...configuration_utils import FrozenDict -from ...image_processor import VaeImageProcessor +from ...image_processor import InpaintProcessor, VaeImageProcessor from ...models import AutoencoderKLQwenImage from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -27,6 +27,24 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def _decode_latents(components: Krea2ModularPipeline, latents: torch.Tensor, height: int, width: int): + vae = components.vae + p = components.patch_size + batch_size, _, channels = latents.shape + latent_height = p * (height // (components.vae_scale_factor * p)) + latent_width = p * (width // (components.vae_scale_factor * p)) + latents = latents.view(batch_size, latent_height // p, latent_width // p, channels // (p * p), p, p) + latents = latents.permute(0, 3, 1, 4, 2, 5) + latents = latents.reshape(batch_size, channels // (p * p), 1, latent_height, latent_width).to(vae.dtype) + + latents_mean = torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1) + latents_std = torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1) + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + latents = latents * latents_std + latents_mean + return vae.decode(latents, return_dict=False)[0][:, :, 0] + + # auto_docstring class Krea2DecodeStep(ModularPipelineBlocks): """ @@ -94,28 +112,82 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + image = _decode_latents(components, block_state.latents, int(block_state.height), int(block_state.width)) + block_state.images = components.image_processor.postprocess(image, output_type=block_state.output_type) - vae = components.vae - p = components.patch_size - latents = block_state.latents + self.set_block_state(state, block_state) + return components, state - batch_size, _, channels = latents.shape - height = p * (int(block_state.height) // (components.vae_scale_factor * p)) - width = p * (int(block_state.width) // (components.vae_scale_factor * p)) - latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p) - latents = latents.permute(0, 3, 1, 4, 2, 5) - latents = latents.reshape(batch_size, channels // (p * p), 1, height, width) - latents = latents.to(vae.dtype) - latents_mean = ( - torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(latents.device, latents.dtype) - ) - latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to( - latents.device, latents.dtype +# auto_docstring +class Krea2InpaintDecodeStep(ModularPipelineBlocks): + """ + Decode Krea 2 inpainting latents and optionally overlay a cropped result on the original image. + + Components: + vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to overlay a cropped inpainting result on the original image. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Decode Krea 2 inpainting latents and optionally overlay a cropped result on the original image." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLQwenImage), + ComponentSpec( + "image_mask_processor", + InpaintProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("output_type", default="pil"), + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + InputParam.template("latents", required=True), + InputParam( + name="mask_overlay_kwargs", + type_hint=dict, + description="Arguments used to overlay a cropped inpainting result on the original image.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("images")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + image = _decode_latents(components, block_state.latents, int(block_state.height), int(block_state.width)) + overlay_kwargs = block_state.mask_overlay_kwargs or {} + block_state.images = components.image_mask_processor.postprocess( + image, output_type=block_state.output_type, **overlay_kwargs ) - latents = latents / latents_std + latents_mean - image = vae.decode(latents, return_dict=False)[0][:, :, 0] - block_state.images = components.image_processor.postprocess(image, output_type=block_state.output_type) self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/krea2/denoise.py b/src/diffusers/modular_pipelines/krea2/denoise.py index 88c6cdca7aba..7ef52191b11e 100644 --- a/src/diffusers/modular_pipelines/krea2/denoise.py +++ b/src/diffusers/modular_pipelines/krea2/denoise.py @@ -208,6 +208,113 @@ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: return components, block_state +class Krea2ReferenceLoopDenoiser(Krea2LoopDenoiser): + model_name = "krea2" + + @property + def description(self) -> str: + return "Run the Krea 2 transformer with clean reference-image tokens prepended to the noisy target tokens." + + @property + def inputs(self) -> list[InputParam]: + return super().inputs + [ + InputParam( + name="reference_image_latents", + required=True, + type_hint=list[torch.Tensor], + description="Packed clean reference-image latents in conditioning order.", + ), + InputParam( + name="reference_attention_scale", + type_hint=float | list[float], + default=1.0, + description="One multiplier for all references or one multiplier per reference in conditioning order.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + transformer = components.transformer + latents = block_state.latents.to(transformer.dtype) + timestep = block_state.timestep.to(transformer.dtype) + reference_image_latents = [latents.to(transformer.dtype) for latents in block_state.reference_image_latents] + guider_inputs = { + "encoder_hidden_states": ( + block_state.prompt_embeds.to(transformer.dtype), + block_state.negative_prompt_embeds.to(transformer.dtype) + if block_state.negative_prompt_embeds is not None + else None, + ), + "encoder_attention_mask": ( + block_state.prompt_embeds_mask, + block_state.negative_prompt_embeds_mask, + ), + } + + components.guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) + guider_state = components.guider.prepare_inputs(guider_inputs) + for guider_state_batch in guider_state: + components.guider.prepare_models(transformer) + cond_kwargs = {name: getattr(guider_state_batch, name) for name in guider_inputs} + guider_state_batch.noise_pred = transformer( + hidden_states=latents, + reference_hidden_states=reference_image_latents, + reference_attention_scale=block_state.reference_attention_scale, + timestep=timestep, + position_ids=block_state.position_ids, + attention_kwargs=block_state.attention_kwargs, + return_dict=False, + **cond_kwargs, + )[0] + components.guider.cleanup_models(transformer) + + block_state.noise_pred = components.guider(guider_state).pred + return components, block_state + + +class Krea2TurboReferenceLoopDenoiser(Krea2TurboLoopDenoiser): + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Run the Krea 2 Turbo transformer with clean reference-image tokens prepended to the noisy target tokens." + ) + + @property + def inputs(self) -> list[InputParam]: + return super().inputs + [ + InputParam( + name="reference_image_latents", + required=True, + type_hint=list[torch.Tensor], + description="Packed clean reference-image latents in conditioning order.", + ), + InputParam( + name="reference_attention_scale", + type_hint=float | list[float], + default=1.0, + description="One multiplier for all references or one multiplier per reference in conditioning order.", + ), + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + transformer = components.transformer + block_state.noise_pred = transformer( + hidden_states=block_state.latents.to(transformer.dtype), + reference_hidden_states=[latents.to(transformer.dtype) for latents in block_state.reference_image_latents], + reference_attention_scale=block_state.reference_attention_scale, + timestep=block_state.timestep.to(transformer.dtype), + position_ids=block_state.position_ids, + attention_kwargs=block_state.attention_kwargs, + encoder_hidden_states=block_state.prompt_embeds.to(transformer.dtype), + encoder_attention_mask=block_state.prompt_embeds_mask, + return_dict=False, + )[0] + return components, block_state + + class Krea2LoopAfterDenoiser(ModularPipelineBlocks): model_name = "krea2" @@ -233,6 +340,43 @@ def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: return components, block_state +class Krea2LoopAfterDenoiserInpaint(ModularPipelineBlocks): + model_name = "krea2" + + @property + def description(self) -> str: + return "Within the denoising loop: preserve the unmasked image latents at the next noise level." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam(name="mask", required=True, type_hint=torch.Tensor, description="The packed inpainting mask."), + InputParam.template("image_latents", required=True), + InputParam( + name="initial_noise", required=True, type_hint=torch.Tensor, description="The sampled initial noise." + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor): + image_latents = block_state.image_latents + if i < len(block_state.timesteps) - 1: + next_timestep = block_state.timesteps[i + 1] + image_latents = components.scheduler.scale_noise( + image_latents, next_timestep.reshape(1), block_state.initial_noise + ) + block_state.latents = (1 - block_state.mask) * image_latents + block_state.mask * block_state.latents + return components, block_state + + class Krea2DenoiseLoopWrapper(LoopSequentialPipelineBlocks): model_name = "krea2" @@ -367,3 +511,203 @@ def description(self) -> str: "latents over `timesteps`, running the transformer on the conditional text features. The distilled " "checkpoint runs without classifier-free guidance." ) + + +# auto_docstring +class Krea2ReferenceDenoiseStep(Krea2DenoiseLoopWrapper): + """ + Denoise Krea 2 target latents while attending to clean reference-image tokens. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + negative_prompt_embeds (`Tensor`, *optional*): + Negative stacked text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative text mask. + reference_image_latents (`list`): + Packed clean reference-image latents in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + + Outputs: + latents (`Tensor`): + The denoised latents. + """ + + model_name = "krea2" + block_classes = [Krea2LoopBeforeDenoiser, Krea2ReferenceLoopDenoiser, Krea2LoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return "Denoise Krea 2 target latents while attending to clean reference-image tokens." + + +# auto_docstring +class Krea2TurboReferenceDenoiseStep(Krea2DenoiseLoopWrapper): + """ + Denoise Krea 2 Turbo target latents while attending to clean reference-image tokens. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + reference_image_latents (`list`): + Packed clean reference-image latents in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + + Outputs: + latents (`Tensor`): + The denoised latents. + """ + + model_name = "krea2" + block_classes = [Krea2LoopBeforeDenoiser, Krea2TurboReferenceLoopDenoiser, Krea2LoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return "Denoise Krea 2 Turbo target latents while attending to clean reference-image tokens." + + +# auto_docstring +class Krea2InpaintDenoiseStep(Krea2DenoiseLoopWrapper): + """ + Krea 2 denoising loop that preserves unmasked source-image latents after every denoising step. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`) transformer + (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + negative_prompt_embeds (`Tensor`, *optional*): + Negative stacked text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Negative text mask. + mask (`Tensor`): + The packed inpainting mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + initial_noise (`Tensor`): + The sampled initial noise. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2LoopBeforeDenoiser, + Krea2LoopDenoiser, + Krea2LoopAfterDenoiser, + Krea2LoopAfterDenoiserInpaint, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser", "inpaint"] + + @property + def description(self) -> str: + return "Krea 2 denoising loop that preserves unmasked source-image latents after every denoising step." + + +# auto_docstring +class Krea2TurboInpaintDenoiseStep(Krea2DenoiseLoopWrapper): + """ + Krea 2 Turbo denoising loop that preserves unmasked source-image latents after every denoising step. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`Krea2Transformer2DModel`) + + Inputs: + timesteps (`Tensor`): + Denoising timesteps from set_timesteps. + num_inference_steps (`int`): + The number of denoising steps. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + latents (`Tensor`): + Packed image latents. + batch_size (`int`): + Effective batch size. + prompt_embeds (`Tensor`): + Conditional stacked text features. + prompt_embeds_mask (`Tensor`): + Conditional text mask. + position_ids (`Tensor`): + Shared rotary coordinates for the [text | image] sequence. + mask (`Tensor`): + The packed inpainting mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + initial_noise (`Tensor`): + The sampled initial noise. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2LoopBeforeDenoiser, + Krea2TurboLoopDenoiser, + Krea2LoopAfterDenoiser, + Krea2LoopAfterDenoiserInpaint, + ] + block_names = ["before_denoiser", "denoiser", "after_denoiser", "inpaint"] + + @property + def description(self) -> str: + return "Krea 2 Turbo denoising loop that preserves unmasked source-image latents after every denoising step." diff --git a/src/diffusers/modular_pipelines/krea2/encoders.py b/src/diffusers/modular_pipelines/krea2/encoders.py index 7640222e9ad2..9d9af2729079 100644 --- a/src/diffusers/modular_pipelines/krea2/encoders.py +++ b/src/diffusers/modular_pipelines/krea2/encoders.py @@ -13,11 +13,14 @@ # limitations under the License. +import PIL.Image import torch -from transformers import AutoTokenizer, Qwen3VLModel +from transformers import AutoTokenizer, Qwen2VLImageProcessor, Qwen3VLModel from ...configuration_utils import FrozenDict from ...guiders import ClassifierFreeGuidance +from ...image_processor import InpaintProcessor, VaeImageProcessor +from ...models import AutoencoderKLQwenImage from ...utils import logging from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam @@ -42,6 +45,42 @@ _PROMPT_TEMPLATE_ENCODE_START_IDX = 34 _PROMPT_TEMPLATE_ENCODE_NUM_SUFFIX_TOKENS = 5 +_REFERENCE_PROMPT_TEMPLATE = ( + "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, " + "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n" + "{}{}<|im_end|>\n<|im_start|>assistant\n" +) + + +class Krea2ReferenceImageProcessor(Qwen2VLImageProcessor): + def __init__( + self, + size: dict | None = None, + patch_size: int = 16, + temporal_patch_size: int = 2, + merge_size: int = 2, + image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5), + image_std: tuple[float, float, float] = (0.5, 0.5, 0.5), + ): + super().__init__( + size=size or {"longest_edge": 16777216, "shortest_edge": 65536}, + patch_size=patch_size, + temporal_patch_size=temporal_patch_size, + merge_size=merge_size, + image_mean=image_mean, + image_std=image_std, + ) + + @property + def device(self): + if self._processor_device is None: + raise AttributeError("Krea2ReferenceImageProcessor is device-independent") + return self._processor_device + + @device.setter + def device(self, value): + self._processor_device = value + # auto_docstring class Krea2TextEncoderStep(ModularPipelineBlocks): @@ -274,3 +313,678 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi self.set_block_state(state, block_state) return components, state + + +# auto_docstring +class Krea2ReferenceTextEncoderStep(ModularPipelineBlocks): + """ + Encode prompts together with a reference image through Qwen3-VL for reference-conditioned Krea 2 generation. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The negative prompt(s) for CFG. + reference_image (`Image | list`): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Encode prompts together with a reference image through Qwen3-VL for reference-conditioned Krea 2 generation." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec( + "reference_image_processor", + Krea2ReferenceImageProcessor, + default_creation_method="from_config", + description="The Qwen3-VL processor used for image-grounded prompt encoding.", + ), + ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."), + ComponentSpec( + "guider", + ClassifierFreeGuidance, + config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam(name="negative_prompt", type_hint=str, description="The negative prompt(s) for CFG."), + InputParam( + name="reference_image", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + required=True, + description="First reference image(s), or scene reference for two-reference generation.", + ), + InputParam( + name="reference_image_2", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + description="Optional second reference image(s), used as the subject reference.", + ), + InputParam( + name="reference_image_encoder_resolution", + type_hint=int, + default=768, + description="Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + OutputParam.template("prompt_embeds_mask"), + OutputParam.template("negative_prompt_embeds"), + OutputParam.template("negative_prompt_embeds_mask"), + ] + + def _encode_prompt(self, components, prompts, reference_images, reference_images_2, encoder_resolution, device): + references_by_input = [] + for name, images in (("reference_image", reference_images), ("reference_image_2", reference_images_2)): + if images is None: + continue + if isinstance(images, PIL.Image.Image): + images = [images] + if len(images) == 1 and len(prompts) > 1: + images = images * len(prompts) + if len(images) != len(prompts): + raise ValueError( + f"`{name}` must contain one image or one image per prompt, but got {len(images)} images for " + f"{len(prompts)} prompts." + ) + references_by_input.append(images) + + processed_images = [] + for prompt_images in zip(*references_by_input): + for image in prompt_images: + image = image.convert("RGB") + if encoder_resolution and max(image.size) > encoder_resolution: + scale = encoder_resolution / max(image.size) + image = image.resize( + (max(16, round(image.width * scale)), max(16, round(image.height * scale))), + PIL.Image.Resampling.LANCZOS, + ) + processed_images.append(image) + + image_inputs = components.reference_image_processor(images=processed_images, return_tensors="pt") + image_token = "<|image_pad|>" + image_token_counts = ( + image_inputs.image_grid_thw.prod(dim=1) // components.reference_image_processor.merge_size**2 + ).tolist() + num_references = len(references_by_input) + vision_block = "<|vision_start|><|image_pad|><|vision_end|>" + texts = [] + for prompt_index, prompt in enumerate(prompts): + prompt_vision_blocks = "" + for reference_index in range(num_references): + count = image_token_counts[prompt_index * num_references + reference_index] + prompt_vision_blocks += vision_block.replace(image_token, image_token * count) + texts.append(_REFERENCE_PROMPT_TEMPLATE.format(prompt_vision_blocks, prompt)) + text_inputs = components.tokenizer(texts, padding=True, return_tensors="pt") + input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + image_token_id = components.tokenizer.convert_tokens_to_ids(image_token) + outputs = components.text_encoder( + input_ids=input_ids, + attention_mask=attention_mask, + pixel_values=image_inputs.pixel_values.to(device), + image_grid_thw=image_inputs.image_grid_thw.to(device), + mm_token_type_ids=input_ids.eq(image_token_id).long(), + output_hidden_states=True, + ) + hidden_states = torch.stack([outputs.hidden_states[i] for i in KREA2_TEXT_ENCODER_SELECT_LAYERS], dim=2) + return ( + hidden_states[:, _PROMPT_TEMPLATE_ENCODE_START_IDX:], + attention_mask[:, _PROMPT_TEMPLATE_ENCODE_START_IDX:].bool(), + ) + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) + + block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( + components, + prompts, + block_state.reference_image, + block_state.reference_image_2, + block_state.reference_image_encoder_resolution, + device, + ) + + block_state.negative_prompt_embeds = None + block_state.negative_prompt_embeds_mask = None + if components.requires_unconditional_embeds: + negative_prompts = block_state.negative_prompt + if negative_prompts is None: + negative_prompts = "" + if isinstance(negative_prompts, str): + negative_prompts = [negative_prompts] * len(prompts) + block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._encode_prompt( + components, + negative_prompts, + block_state.reference_image, + block_state.reference_image_2, + block_state.reference_image_encoder_resolution, + device, + ) + prompt_length = block_state.prompt_embeds.shape[1] + negative_length = block_state.negative_prompt_embeds.shape[1] + if prompt_length < negative_length: + padding = negative_length - prompt_length + block_state.prompt_embeds = torch.nn.functional.pad( + block_state.prompt_embeds, (0, 0, 0, 0, 0, padding) + ) + block_state.prompt_embeds_mask = torch.nn.functional.pad( + block_state.prompt_embeds_mask, (0, padding), value=False + ) + elif negative_length < prompt_length: + padding = prompt_length - negative_length + block_state.negative_prompt_embeds = torch.nn.functional.pad( + block_state.negative_prompt_embeds, (0, 0, 0, 0, 0, padding) + ) + block_state.negative_prompt_embeds_mask = torch.nn.functional.pad( + block_state.negative_prompt_embeds_mask, (0, padding), value=False + ) + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2TurboReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): + """ + Encode prompts with a reference image for reference-conditioned Krea 2 Turbo generation. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + reference_image (`Image | list`): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + """ + + @property + def description(self) -> str: + return "Encode prompts with a reference image for reference-conditioned Krea 2 Turbo generation." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", Qwen3VLModel, description="The Qwen3-VL text encoder."), + ComponentSpec( + "reference_image_processor", + Krea2ReferenceImageProcessor, + default_creation_method="from_config", + description="The Qwen3-VL processor used for image-grounded prompt encoding.", + ), + ComponentSpec("tokenizer", AutoTokenizer, description="The tokenizer paired with the text encoder."), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("prompt", required=True), + InputParam( + name="reference_image", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + required=True, + description="First reference image(s), or scene reference for two-reference generation.", + ), + InputParam( + name="reference_image_2", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + description="Optional second reference image(s), used as the subject reference.", + ), + InputParam( + name="reference_image_encoder_resolution", + type_hint=int, + default=768, + description="Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("prompt_embeds"), OutputParam.template("prompt_embeds_mask")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + prompts = [block_state.prompt] if isinstance(block_state.prompt, str) else list(block_state.prompt) + block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt( + components, + prompts, + block_state.reference_image, + block_state.reference_image_2, + block_state.reference_image_encoder_resolution, + components._execution_device, + ) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2ProcessImagesInputStep(ModularPipelineBlocks): + """ + Preprocess an input image for Krea 2 image-to-image generation. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + processed_image (`Tensor`): + The preprocessed input image. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Preprocess an input image for Krea 2 image-to-image generation." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", required=True), + InputParam.template("height"), + InputParam.template("width"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="processed_image", type_hint=torch.Tensor, description="The preprocessed input image.") + ] + + @staticmethod + def check_inputs(height, width, multiple): + if height is not None and height % multiple != 0: + raise ValueError(f"`height` must be divisible by {multiple}, but is {height}") + if width is not None and width % multiple != 0: + raise ValueError(f"`width` must be divisible by {multiple}, but is {width}") + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs(block_state.height, block_state.width, components.image_processor.config.vae_scale_factor) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + block_state.processed_image = components.image_processor.preprocess( + image=block_state.image, height=height, width=width + ) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2InpaintProcessImagesInputStep(ModularPipelineBlocks): + """ + Preprocess an input image and mask together for Krea 2 inpainting. + + Components: + image_mask_processor (`InpaintProcessor`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_image (`Tensor`): + The preprocessed input image. + processed_mask_image (`Tensor`): + The preprocessed inpainting mask. + mask_overlay_kwargs (`dict`): + Arguments used to overlay a cropped inpainting result on the original image. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Preprocess an input image and mask together for Krea 2 inpainting." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_mask_processor", + InpaintProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("image", required=True), + InputParam.template("mask_image", required=True), + InputParam.template("height"), + InputParam.template("width"), + InputParam.template("padding_mask_crop"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam(name="processed_image", type_hint=torch.Tensor, description="The preprocessed input image."), + OutputParam( + name="processed_mask_image", type_hint=torch.Tensor, description="The preprocessed inpainting mask." + ), + OutputParam( + name="mask_overlay_kwargs", + type_hint=dict, + description="Arguments used to overlay a cropped inpainting result on the original image.", + ), + ] + + @staticmethod + def check_inputs(height, width, multiple): + if height is not None and height % multiple != 0: + raise ValueError(f"`height` must be divisible by {multiple}, but is {height}") + if width is not None and width % multiple != 0: + raise ValueError(f"`width` must be divisible by {multiple}, but is {width}") + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + self.check_inputs( + block_state.height, block_state.width, components.image_mask_processor.config.vae_scale_factor + ) + height = block_state.height or components.default_height + width = block_state.width or components.default_width + block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = ( + components.image_mask_processor.preprocess( + image=block_state.image, + mask=block_state.mask_image, + height=height, + width=width, + padding_mask_crop=block_state.padding_mask_crop, + ) + ) + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2VaeEncoderStep(ModularPipelineBlocks): + """ + Encode a preprocessed image into normalized Krea 2 image latents. + + Components: + vae (`AutoencoderKLQwenImage`) + + Inputs: + processed_image (`Tensor`): + The preprocessed image. + + Outputs: + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Encode a preprocessed image into normalized Krea 2 image latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLQwenImage)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_image", required=True, type_hint=torch.Tensor, description="The preprocessed image." + ) + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam.template("image_latents")] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + image = block_state.processed_image + if image.ndim == 4: + image = image.unsqueeze(2) + elif image.ndim != 5: + raise ValueError(f"`processed_image` must have 4 or 5 dimensions, but got {image.ndim}") + + image = image.to(device=components._execution_device, dtype=components.vae.dtype) + image_latents = components.vae.encode(image).latent_dist.mode() + + latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, components.vae.config.z_dim, 1, 1, 1) + latents_std = torch.tensor(components.vae.config.latents_std).view(1, components.vae.config.z_dim, 1, 1, 1) + latents_mean = latents_mean.to(image_latents.device, image_latents.dtype) + latents_std = latents_std.to(image_latents.device, image_latents.dtype) + block_state.image_latents = (image_latents - latents_mean) / latents_std + + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2ReferenceProcessImagesInputStep(ModularPipelineBlocks): + """ + Preprocess a reference image at the target output resolution for VAE encoding. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + reference_image (`Image | list`): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + + Outputs: + processed_reference_images (`list`): + Reference images resized and normalized for VAE encoding in conditioning order. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Preprocess a reference image at the target output resolution for VAE encoding." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "image_processor", + VaeImageProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ) + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="reference_image", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + required=True, + description="First reference image(s), or scene reference for two-reference generation.", + ), + InputParam( + name="reference_image_2", + type_hint=PIL.Image.Image | list[PIL.Image.Image], + description="Optional second reference image(s), used as the subject reference.", + ), + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="processed_reference_images", + type_hint=list[torch.Tensor], + description="Reference images resized and normalized for VAE encoding in conditioning order.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + multiple = components.image_processor.config.vae_scale_factor + if block_state.height % multiple != 0 or block_state.width % multiple != 0: + raise ValueError(f"`height` and `width` must be divisible by {multiple} for reference conditioning.") + reference_images = [block_state.reference_image] + if block_state.reference_image_2 is not None: + reference_images.append(block_state.reference_image_2) + block_state.processed_reference_images = [ + components.image_processor.preprocess(image=image, height=block_state.height, width=block_state.width) + for image in reference_images + ] + self.set_block_state(state, block_state) + return components, state + + +# auto_docstring +class Krea2ReferenceVaeEncoderStep(ModularPipelineBlocks): + """ + Encode a preprocessed reference image into normalized Krea 2 latents. + + Components: + vae (`AutoencoderKLQwenImage`) + + Inputs: + processed_reference_images (`list`): + The preprocessed reference images in conditioning order. + + Outputs: + reference_image_latents (`list`): + Normalized latent representations of the reference images in conditioning order. + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return "Encode a preprocessed reference image into normalized Krea 2 latents." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ComponentSpec("vae", AutoencoderKLQwenImage)] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="processed_reference_images", + type_hint=list[torch.Tensor], + required=True, + description="The preprocessed reference images in conditioning order.", + ) + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + name="reference_image_latents", + type_hint=list[torch.Tensor], + description="Normalized latent representations of the reference images in conditioning order.", + ) + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latents_mean = torch.tensor(components.vae.config.latents_mean).view(1, components.vae.config.z_dim, 1, 1, 1) + latents_std = torch.tensor(components.vae.config.latents_std).view(1, components.vae.config.z_dim, 1, 1, 1) + block_state.reference_image_latents = [] + for processed_reference_image in block_state.processed_reference_images: + reference_image = processed_reference_image.unsqueeze(2).to( + device=components._execution_device, dtype=components.vae.dtype + ) + reference_image_latents = components.vae.encode(reference_image).latent_dist.mode() + reference_image_latents = ( + reference_image_latents - latents_mean.to(reference_image_latents) + ) / latents_std.to(reference_image_latents) + block_state.reference_image_latents.append(reference_image_latents) + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py index ae3b2ac4fb52..0e067b40b02d 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py +++ b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py @@ -14,22 +14,81 @@ from ...utils import logging -from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam from .before_denoise import ( + Krea2ApplyStrengthStep, + Krea2ImageInputsStep, + Krea2PrepareImageLatentsStep, Krea2PrepareLatentsStep, + Krea2PrepareMaskLatentsStep, Krea2PreparePositionIdsStep, + Krea2PrepareReferencePositionIdsStep, + Krea2ReferenceInputsStep, Krea2SetTimestepsStep, Krea2TextInputsStep, ) -from .decoders import Krea2DecodeStep -from .denoise import Krea2DenoiseStep -from .encoders import Krea2TextEncoderStep +from .decoders import Krea2DecodeStep, Krea2InpaintDecodeStep +from .denoise import Krea2DenoiseStep, Krea2InpaintDenoiseStep, Krea2ReferenceDenoiseStep +from .encoders import ( + Krea2InpaintProcessImagesInputStep, + Krea2ProcessImagesInputStep, + Krea2ReferenceProcessImagesInputStep, + Krea2ReferenceTextEncoderStep, + Krea2ReferenceVaeEncoderStep, + Krea2TextEncoderStep, + Krea2VaeEncoderStep, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# auto_docstring +class Krea2AutoTextEncoderStep(AutoPipelineBlocks): + """ + Select text-only or reference-image-grounded Krea 2 prompt encoding. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. guider (`ClassifierFreeGuidance`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The negative prompt(s) for CFG. + reference_image (`Image | list`, *optional*): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + negative_prompt_embeds (`Tensor`): + The negative prompt embeddings. + negative_prompt_embeds_mask (`Tensor`): + The negative prompt embeddings mask. + """ + + model_name = "krea2" + block_classes = [Krea2ReferenceTextEncoderStep, Krea2TextEncoderStep] + block_names = ["reference", "text"] + block_trigger_inputs = ["reference_image", None] + + @property + def description(self) -> str: + return "Select text-only or reference-image-grounded Krea 2 prompt encoding." + + CORE_DENOISE_BLOCKS = InsertableDict( [ ("input", Krea2TextInputsStep()), @@ -101,45 +160,701 @@ def outputs(self) -> list[OutputParam]: ] +# auto_docstring +class Krea2Img2ImgVaeEncoderStep(SequentialPipelineBlocks): + """ + Preprocess and VAE-encode an image for Krea 2 image-to-image generation. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + processed_image (`Tensor`): + The preprocessed input image. + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "krea2" + block_classes = [Krea2ProcessImagesInputStep, Krea2VaeEncoderStep] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Preprocess and VAE-encode an image for Krea 2 image-to-image generation." + + +# auto_docstring +class Krea2InpaintVaeEncoderStep(SequentialPipelineBlocks): + """ + Preprocess an image and mask and VAE-encode the image for Krea 2 inpainting. + + Components: + image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLQwenImage`) + + Inputs: + image (`Image | list`): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`): + Mask image for inpainting. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_image (`Tensor`): + The preprocessed input image. + processed_mask_image (`Tensor`): + The preprocessed inpainting mask. + mask_overlay_kwargs (`dict`): + Arguments used to overlay a cropped inpainting result on the original image. + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "krea2" + block_classes = [Krea2InpaintProcessImagesInputStep, Krea2VaeEncoderStep] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Preprocess an image and mask and VAE-encode the image for Krea 2 inpainting." + + +# auto_docstring +class Krea2ReferenceVaeEncoderBlocks(SequentialPipelineBlocks): + """ + Preprocess and VAE-encode an image for reference-conditioned Krea 2 generation. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) + + Inputs: + reference_image (`Image | list`): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + + Outputs: + processed_reference_images (`list`): + Reference images resized and normalized for VAE encoding in conditioning order. + reference_image_latents (`list`): + Normalized latent representations of the reference images in conditioning order. + """ + + model_name = "krea2" + block_classes = [Krea2ReferenceProcessImagesInputStep, Krea2ReferenceVaeEncoderStep] + block_names = ["preprocess", "encode"] + + @property + def description(self) -> str: + return "Preprocess and VAE-encode an image for reference-conditioned Krea 2 generation." + + +# auto_docstring +class Krea2AutoVaeEncoderStep(AutoPipelineBlocks): + """ + Select the Krea 2 inpainting or image-to-image VAE encoder from the provided image inputs. + + Components: + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) image_mask_processor + (`InpaintProcessor`) + + Inputs: + reference_image (`Image | list`, *optional*): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The width in pixels of the generated image. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`, *optional*): + Mask image for inpainting. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. + + Outputs: + processed_reference_images (`list`): + Reference images resized and normalized for VAE encoding in conditioning order. + reference_image_latents (`list`): + Normalized latent representations of the reference images in conditioning order. + processed_image (`Tensor`): + The preprocessed input image. + processed_mask_image (`Tensor`): + The preprocessed inpainting mask. + mask_overlay_kwargs (`dict`): + Arguments used to overlay a cropped inpainting result on the original image. + image_latents (`Tensor`): + The latent representation of the input image. + """ + + model_name = "krea2" + block_classes = [Krea2ReferenceVaeEncoderBlocks, Krea2InpaintVaeEncoderStep, Krea2Img2ImgVaeEncoderStep] + block_names = ["reference", "inpaint", "img2img"] + block_trigger_inputs = ["reference_image", "mask_image", "image"] + + @property + def description(self) -> str: + return "Select the Krea 2 inpainting or image-to-image VAE encoder from the provided image inputs." + + +# auto_docstring +class Krea2Img2ImgInputStep(SequentialPipelineBlocks): + """ + Expand Krea 2 text and source-image inputs to the effective image-to-image batch. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + negative_prompt_embeds (`Tensor`): + Negative text features, batch-expanded. + negative_prompt_embeds_mask (`Tensor`): + Negative text mask, batch-expanded. + image_latents (`Tensor`): + The latent representation of the input image. + processed_mask_image (`Tensor`): + The batch-expanded inpainting mask. + height (`int`): + The generation height inferred from the image. + width (`int`): + The generation width inferred from the image. + """ + + model_name = "krea2" + block_classes = [Krea2TextInputsStep, Krea2ImageInputsStep] + block_names = ["text_inputs", "image_inputs"] + + @property + def description(self) -> str: + return "Expand Krea 2 text and source-image inputs to the effective image-to-image batch." + + +# auto_docstring +class Krea2InpaintPrepareLatentsStep(SequentialPipelineBlocks): + """ + Add noise to Krea 2 source-image latents and prepare packed mask latents for inpainting. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + latents (`Tensor`): + Pre-generated noisy latents for image generation. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + timesteps (`Tensor`): + The selected denoising timesteps. + processed_mask_image (`Tensor`): + The preprocessed inpainting mask. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + dtype (`dtype`, *optional*, defaults to torch.float32): + The dtype of the model inputs, can be generated in input step. + + Outputs: + initial_noise (`Tensor`): + The sampled initial noise. + latents (`Tensor`): + Denoised latents. + mask (`Tensor`): + The packed latent-space mask. + """ + + model_name = "krea2" + block_classes = [Krea2PrepareImageLatentsStep, Krea2PrepareMaskLatentsStep] + block_names = ["add_noise", "prepare_mask"] + + @property + def description(self) -> str: + return "Add noise to Krea 2 source-image latents and prepare packed mask latents for inpainting." + + +# auto_docstring +class Krea2ReferenceInputStep(SequentialPipelineBlocks): + """ + Expand Krea 2 text and reference-image conditioning to the effective batch. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + reference_image_latents (`list`): + Normalized reference-image latents from the VAE encoder in conditioning order. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + negative_prompt_embeds (`Tensor`): + Negative text features, batch-expanded. + negative_prompt_embeds_mask (`Tensor`): + Negative text mask, batch-expanded. + reference_image_latents (`list`): + Packed reference-image latents expanded to the effective batch in conditioning order. + """ + + model_name = "krea2" + block_classes = [Krea2TextInputsStep, Krea2ReferenceInputsStep] + block_names = ["text_inputs", "reference_inputs"] + + @property + def description(self) -> str: + return "Expand Krea 2 text and reference-image conditioning to the effective batch." + + +# auto_docstring +class Krea2ReferenceCoreDenoiseStep(SequentialPipelineBlocks): + """ + Generate Krea 2 target latents from noise while attending to clean reference-image tokens. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + reference_image_latents (`list`): + Normalized reference-image latents from the VAE encoder in conditioning order. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2ReferenceInputStep, + Krea2PrepareLatentsStep, + Krea2SetTimestepsStep, + Krea2PrepareReferencePositionIdsStep, + Krea2ReferenceDenoiseStep, + ] + block_names = ["input", "prepare_latents", "set_timesteps", "prepare_position_ids", "denoise"] + + @property + def description(self) -> str: + return "Generate Krea 2 target latents from noise while attending to clean reference-image tokens." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2Img2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Krea 2 image-to-image workflow with strength-adjusted flow-matching denoising. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2Img2ImgInputStep, + Krea2PrepareLatentsStep, + Krea2SetTimestepsStep, + Krea2ApplyStrengthStep, + Krea2PrepareImageLatentsStep, + Krea2PreparePositionIdsStep, + Krea2DenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "apply_strength", + "prepare_image_latents", + "prepare_position_ids", + "denoise", + ] + + @property + def description(self) -> str: + return "Core Krea 2 image-to-image workflow with strength-adjusted flow-matching denoising." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2InpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Krea 2 inpainting workflow with masked latent blending after every denoising step. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 28): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2Img2ImgInputStep, + Krea2PrepareLatentsStep, + Krea2SetTimestepsStep, + Krea2ApplyStrengthStep, + Krea2InpaintPrepareLatentsStep, + Krea2PreparePositionIdsStep, + Krea2InpaintDenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "apply_strength", + "prepare_inpaint_latents", + "prepare_position_ids", + "denoise", + ] + + @property + def description(self) -> str: + return "Core Krea 2 inpainting workflow with masked latent blending after every denoising step." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2AutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Select the Krea 2 text-to-image, image-to-image, or inpainting denoising workflow. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`ClassifierFreeGuidance`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + negative_prompt_embeds (`Tensor`, *optional*): + Per-prompt negative text features. + negative_prompt_embeds_mask (`Tensor`, *optional*): + Per-prompt negative text mask. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + reference_image_latents (`list`, *optional*): + Normalized reference-image latents from the VAE encoder in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2CoreDenoiseStep, + Krea2ReferenceCoreDenoiseStep, + Krea2InpaintCoreDenoiseStep, + Krea2Img2ImgCoreDenoiseStep, + ] + block_names = ["text2image", "reference", "inpaint", "img2img"] + block_trigger_inputs = ["reference_image_latents", "processed_mask_image", "image_latents"] + default_block_name = "text2image" + + def select_block(self, reference_image_latents=None, processed_mask_image=None, image_latents=None): + if reference_image_latents is not None: + return "reference" + if processed_mask_image is not None: + return "inpaint" + if image_latents is not None: + return "img2img" + return "text2image" + + @property + def description(self) -> str: + return "Select the Krea 2 text-to-image, image-to-image, or inpainting denoising workflow." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2AutoDecodeStep(AutoPipelineBlocks): + """ + Select the standard or inpainting-aware Krea 2 decoder. + + Components: + vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`) image_processor + (`VaeImageProcessor`) + + Inputs: + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + latents (`Tensor`): + Pre-generated noisy latents for image generation. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to overlay a cropped inpainting result on the original image. + + Outputs: + images (`list`): + Generated images. + """ + + model_name = "krea2" + block_classes = [Krea2InpaintDecodeStep, Krea2DecodeStep] + block_names = ["inpaint", "default"] + block_trigger_inputs = ["mask", None] + + @property + def description(self) -> str: + return "Select the standard or inpainting-aware Krea 2 decoder." + + # auto_docstring class Krea2AutoBlocks(SequentialPipelineBlocks): """ - Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise (symmetric CFG) -> - decode. + Auto Modular pipeline for text-to-image, image-to-image, and inpainting using Krea 2. Supported workflows: - `text2image`: requires `prompt` + - `image2image`: requires `prompt`, `image` + - `inpainting`: requires `prompt`, `image`, `mask_image` + - `reference`: requires `prompt`, `reference_image` Components: - text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired - with the text encoder. guider (`ClassifierFreeGuidance`) transformer (`Krea2Transformer2DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. guider (`ClassifierFreeGuidance`) + image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`) image_mask_processor + (`InpaintProcessor`) transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: prompt (`str`): The prompt or prompts to guide image generation. negative_prompt (`str`, *optional*): The negative prompt(s) for CFG. + reference_image (`Image | list`, *optional*): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): Maximum sequence length for prompt encoding. + height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The width in pixels of the generated image. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`, *optional*): + Mask image for inpainting. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - latents (`Tensor`, *optional*): + latents (`Tensor`): Pre-generated noisy latents for image generation. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`, *optional*, defaults to 28): + num_inference_steps (`int`): The number of denoising steps. sigmas (`list`, *optional*): Custom sigma schedule (defaults to a linear ramp). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. + reference_image_latents (`list`, *optional*): + Normalized reference-image latents from the VAE encoder in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to overlay a cropped inpainting result on the original image. Outputs: images (`list`): @@ -148,22 +863,23 @@ class Krea2AutoBlocks(SequentialPipelineBlocks): model_name = "krea2" block_classes = [ - Krea2TextEncoderStep, - Krea2CoreDenoiseStep, - Krea2DecodeStep, + Krea2AutoTextEncoderStep, + Krea2AutoVaeEncoderStep, + Krea2AutoCoreDenoiseStep, + Krea2AutoDecodeStep, ] - block_names = ["text_encoder", "denoise", "decode"] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] _workflow_map = { "text2image": {"prompt": True}, + "image2image": {"prompt": True, "image": True}, + "inpainting": {"prompt": True, "image": True, "mask_image": True}, + "reference": {"prompt": True, "reference_image": True}, } @property def description(self) -> str: - return ( - "Auto Modular pipeline for text-to-image generation using Krea 2: encode text -> core denoise " - "(symmetric CFG) -> decode." - ) + return "Auto Modular pipeline for text-to-image, image-to-image, and inpainting using Krea 2." @property def outputs(self) -> list[OutputParam]: diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py index 79fa5406c4e5..3568b1430a18 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py +++ b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py @@ -14,22 +14,70 @@ from ...utils import logging -from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict, OutputParam from .before_denoise import ( + Krea2ApplyStrengthStep, + Krea2ImageInputsStep, + Krea2PrepareImageLatentsStep, Krea2PrepareLatentsStep, Krea2PreparePositionIdsStep, + Krea2PrepareReferencePositionIdsStep, + Krea2ReferenceInputsStep, Krea2TurboSetTimestepsStep, Krea2TurboTextInputsStep, ) -from .decoders import Krea2DecodeStep -from .denoise import Krea2TurboDenoiseStep -from .encoders import Krea2TurboTextEncoderStep +from .denoise import Krea2TurboDenoiseStep, Krea2TurboInpaintDenoiseStep, Krea2TurboReferenceDenoiseStep +from .encoders import Krea2TurboReferenceTextEncoderStep, Krea2TurboTextEncoderStep +from .modular_blocks_krea2 import ( + Krea2AutoDecodeStep, + Krea2AutoVaeEncoderStep, + Krea2InpaintPrepareLatentsStep, +) logger = logging.get_logger(__name__) # pylint: disable=invalid-name +# auto_docstring +class Krea2TurboAutoTextEncoderStep(AutoPipelineBlocks): + """ + Select text-only or reference-image-grounded Krea 2 Turbo prompt encoding. + + Components: + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + reference_image (`Image | list`, *optional*): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. + max_sequence_length (`int`, *optional*, defaults to 512): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + """ + + model_name = "krea2" + block_classes = [Krea2TurboReferenceTextEncoderStep, Krea2TurboTextEncoderStep] + block_names = ["reference", "text"] + block_trigger_inputs = ["reference_image", None] + + @property + def description(self) -> str: + return "Select text-only or reference-image-grounded Krea 2 Turbo prompt encoding." + + CORE_DENOISE_BLOCKS = InsertableDict( [ ("input", Krea2TurboTextInputsStep()), @@ -97,43 +145,430 @@ def outputs(self) -> list[OutputParam]: ] +# auto_docstring +class Krea2TurboImg2ImgInputStep(SequentialPipelineBlocks): + """ + Expand Krea 2 Turbo text and source-image inputs to the effective image-to-image batch. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + image_latents (`Tensor`): + The latent representation of the input image. + processed_mask_image (`Tensor`): + The batch-expanded inpainting mask. + height (`int`): + The generation height inferred from the image. + width (`int`): + The generation width inferred from the image. + """ + + model_name = "krea2" + block_classes = [Krea2TurboTextInputsStep, Krea2ImageInputsStep] + block_names = ["text_inputs", "image_inputs"] + + @property + def description(self) -> str: + return "Expand Krea 2 Turbo text and source-image inputs to the effective image-to-image batch." + + +# auto_docstring +class Krea2TurboReferenceInputStep(SequentialPipelineBlocks): + """ + Expand Krea 2 Turbo text and reference-image conditioning to the effective batch. + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + reference_image_latents (`list`): + Normalized reference-image latents from the VAE encoder in conditioning order. + + Outputs: + batch_size (`int`): + Effective batch size (num prompts * num_images_per_prompt). + dtype (`dtype`): + The dtype of the text features. + prompt_embeds (`Tensor`): + Text features, batch-expanded. + prompt_embeds_mask (`Tensor`): + Text mask, batch-expanded. + reference_image_latents (`list`): + Packed reference-image latents expanded to the effective batch in conditioning order. + """ + + model_name = "krea2" + block_classes = [Krea2TurboTextInputsStep, Krea2ReferenceInputsStep] + block_names = ["text_inputs", "reference_inputs"] + + @property + def description(self) -> str: + return "Expand Krea 2 Turbo text and reference-image conditioning to the effective batch." + + +# auto_docstring +class Krea2TurboReferenceCoreDenoiseStep(SequentialPipelineBlocks): + """ + Generate Krea 2 Turbo target latents from noise while attending to clean reference-image tokens. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + reference_image_latents (`list`): + Normalized reference-image latents from the VAE encoder in conditioning order. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2TurboReferenceInputStep, + Krea2PrepareLatentsStep, + Krea2TurboSetTimestepsStep, + Krea2PrepareReferencePositionIdsStep, + Krea2TurboReferenceDenoiseStep, + ] + block_names = ["input", "prepare_latents", "set_timesteps", "prepare_position_ids", "denoise"] + + @property + def description(self) -> str: + return "Generate Krea 2 Turbo target latents from noise while attending to clean reference-image tokens." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2TurboImg2ImgCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Krea 2 Turbo image-to-image workflow with strength-adjusted denoising. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2TurboImg2ImgInputStep, + Krea2PrepareLatentsStep, + Krea2TurboSetTimestepsStep, + Krea2ApplyStrengthStep, + Krea2PrepareImageLatentsStep, + Krea2PreparePositionIdsStep, + Krea2TurboDenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "apply_strength", + "prepare_image_latents", + "prepare_position_ids", + "denoise", + ] + + @property + def description(self) -> str: + return "Core Krea 2 Turbo image-to-image workflow with strength-adjusted denoising." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2TurboInpaintCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core Krea 2 Turbo inpainting workflow with masked latent blending after every denoising step. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + image_latents (`Tensor`): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + height (`int`, *optional*): + The height in pixels of the generated image. + width (`int`, *optional*): + The width in pixels of the generated image. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 8): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2TurboImg2ImgInputStep, + Krea2PrepareLatentsStep, + Krea2TurboSetTimestepsStep, + Krea2ApplyStrengthStep, + Krea2InpaintPrepareLatentsStep, + Krea2PreparePositionIdsStep, + Krea2TurboInpaintDenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "apply_strength", + "prepare_inpaint_latents", + "prepare_position_ids", + "denoise", + ] + + @property + def description(self) -> str: + return "Core Krea 2 Turbo inpainting workflow with masked latent blending after every denoising step." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + +# auto_docstring +class Krea2TurboAutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Select the Krea 2 Turbo text-to-image, image-to-image, or inpainting denoising workflow. + + Components: + transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + + Inputs: + num_images_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + prompt_embeds (`Tensor`): + Per-prompt stacked text features (B, text_seq_len, num_text_layers, text_hidden_dim). + prompt_embeds_mask (`Tensor`): + Per-prompt boolean text mask (B, text_seq_len). + latents (`Tensor`): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The width in pixels of the generated image. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`): + The number of denoising steps. + sigmas (`list`, *optional*): + Custom sigma schedule (defaults to a linear ramp). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + reference_image_latents (`list`, *optional*): + Normalized reference-image latents from the VAE encoder in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. + + Outputs: + latents (`Tensor`): + Denoised latents. + """ + + model_name = "krea2" + block_classes = [ + Krea2TurboCoreDenoiseStep, + Krea2TurboReferenceCoreDenoiseStep, + Krea2TurboInpaintCoreDenoiseStep, + Krea2TurboImg2ImgCoreDenoiseStep, + ] + block_names = ["text2image", "reference", "inpaint", "img2img"] + block_trigger_inputs = ["reference_image_latents", "processed_mask_image", "image_latents"] + default_block_name = "text2image" + + def select_block(self, reference_image_latents=None, processed_mask_image=None, image_latents=None): + if reference_image_latents is not None: + return "reference" + if processed_mask_image is not None: + return "inpaint" + if image_latents is not None: + return "img2img" + return "text2image" + + @property + def description(self) -> str: + return "Select the Krea 2 Turbo text-to-image, image-to-image, or inpainting denoising workflow." + + @property + def outputs(self) -> list[OutputParam]: + return [OutputParam.template("latents")] + + # auto_docstring class Krea2TurboAutoBlocks(SequentialPipelineBlocks): """ - Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode text -> core - denoise (guidance-free) -> decode. + Auto Modular pipeline for text-to-image, image-to-image, and inpainting using the distilled Krea 2 Turbo + checkpoint. Supported workflows: - `text2image`: requires `prompt` + - `image2image`: requires `prompt`, `image` + - `inpainting`: requires `prompt`, `image`, `mask_image` + - `reference`: requires `prompt`, `reference_image` Components: - text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. tokenizer (`AutoTokenizer`): The tokenizer paired - with the text encoder. transformer (`Krea2Transformer2DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) - vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) + text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor + (`Krea2ReferenceImageProcessor`): The Qwen3-VL processor used for image-grounded prompt encoding. tokenizer + (`AutoTokenizer`): The tokenizer paired with the text encoder. image_processor (`VaeImageProcessor`) vae + (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`) transformer (`Krea2Transformer2DModel`) + scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: prompt (`str`): The prompt or prompts to guide image generation. + reference_image (`Image | list`, *optional*): + First reference image(s), or scene reference for two-reference generation. + reference_image_2 (`Image | list`, *optional*): + Optional second reference image(s), used as the subject reference. + reference_image_encoder_resolution (`int`, *optional*, defaults to 768): + Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): Maximum sequence length for prompt encoding. + height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): + The width in pixels of the generated image. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + mask_image (`Image`, *optional*): + Mask image for inpainting. + padding_mask_crop (`int`, *optional*): + Padding for mask cropping in inpainting. num_images_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - latents (`Tensor`, *optional*): + latents (`Tensor`): Pre-generated noisy latents for image generation. - height (`int`, *optional*, defaults to 1024): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 1024): - The width in pixels of the generated image. generator (`Generator`, *optional*): Torch generator for deterministic generation. - num_inference_steps (`int`, *optional*, defaults to 8): + num_inference_steps (`int`): The number of denoising steps. sigmas (`list`, *optional*): Custom sigma schedule (defaults to a linear ramp). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. + reference_image_latents (`list`, *optional*): + Normalized reference-image latents from the VAE encoder in conditioning order. + reference_attention_scale (`float | list`, *optional*, defaults to 1.0): + One multiplier for all references or one multiplier per reference in conditioning order. + image_latents (`Tensor`, *optional*): + image latents used to guide the image generation. Can be generated from vae_encoder step. + processed_mask_image (`Tensor`, *optional*): + The preprocessed inpainting mask. + strength (`float`, *optional*, defaults to 0.9): + Strength for img2img/inpainting. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. + mask_overlay_kwargs (`dict`, *optional*): + Arguments used to overlay a cropped inpainting result on the original image. Outputs: images (`list`): @@ -142,21 +577,25 @@ class Krea2TurboAutoBlocks(SequentialPipelineBlocks): model_name = "krea2" block_classes = [ - Krea2TurboTextEncoderStep, - Krea2TurboCoreDenoiseStep, - Krea2DecodeStep, + Krea2TurboAutoTextEncoderStep, + Krea2AutoVaeEncoderStep, + Krea2TurboAutoCoreDenoiseStep, + Krea2AutoDecodeStep, ] - block_names = ["text_encoder", "denoise", "decode"] + block_names = ["text_encoder", "vae_encoder", "denoise", "decode"] _workflow_map = { "text2image": {"prompt": True}, + "image2image": {"prompt": True, "image": True}, + "inpainting": {"prompt": True, "image": True, "mask_image": True}, + "reference": {"prompt": True, "reference_image": True}, } @property def description(self) -> str: return ( - "Auto Modular pipeline for text-to-image generation using the distilled Krea 2 turbo checkpoint: encode " - "text -> core denoise (guidance-free) -> decode." + "Auto Modular pipeline for text-to-image, image-to-image, and inpainting using the distilled Krea 2 " + "Turbo checkpoint." ) @property diff --git a/tests/models/transformers/test_models_transformer_krea2.py b/tests/models/transformers/test_models_transformer_krea2.py index 265bc42888ef..6b5cea8ff27f 100644 --- a/tests/models/transformers/test_models_transformer_krea2.py +++ b/tests/models/transformers/test_models_transformer_krea2.py @@ -131,6 +131,37 @@ def get_dummy_inputs(self, height: int | None = None, width: int | None = None) class TestKrea2TransformerModel(Krea2TransformerTesterConfig, ModelTesterMixin): """Core model tests for the Krea 2 Transformer.""" + def test_reference_hidden_states(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + inputs = self.get_dummy_inputs() + image_seq_len = inputs["hidden_states"].shape[1] + text_seq_len = inputs["encoder_hidden_states"].shape[1] + inputs["reference_hidden_states"] = [ + randn_tensor( + inputs["hidden_states"].shape, + generator=self.generator, + device=torch_device, + dtype=self.torch_dtype, + ) + for _ in range(2) + ] + position_ids = torch.cat( + [ + inputs["position_ids"][:text_seq_len], + inputs["position_ids"][text_seq_len:].clone(), + inputs["position_ids"][text_seq_len:].clone(), + inputs["position_ids"][text_seq_len:], + ] + ) + position_ids[text_seq_len : text_seq_len + image_seq_len, 0] = 1 + position_ids[text_seq_len + image_seq_len : text_seq_len + 2 * image_seq_len, 0] = 2 + inputs["position_ids"] = position_ids + + output = model(**inputs, reference_attention_scale=[1.0, 1.0]).sample + boosted_output = model(**inputs, reference_attention_scale=[1.0, 2.0]).sample + assert output.shape == inputs["hidden_states"].shape + assert not torch.allclose(output, boosted_output) + class TestKrea2TransformerMemory(Krea2TransformerTesterConfig, MemoryTesterMixin): """Memory optimization tests for the Krea 2 Transformer.""" diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index 15caa5abe45f..b8cd6c45ea08 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -14,6 +14,9 @@ # limitations under the License. +import PIL +import torch + from diffusers.modular_pipelines import Krea2AutoBlocks, Krea2ModularPipeline from ..test_modular_pipelines_common import ModularPipelineTesterMixin @@ -29,6 +32,47 @@ ("denoise.denoise", "Krea2DenoiseStep"), ("decode", "Krea2DecodeStep"), ], + "image2image": [ + ("text_encoder", "Krea2TextEncoderStep"), + ("vae_encoder.preprocess", "Krea2ProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2VaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TextInputsStep"), + ("denoise.input.image_inputs", "Krea2ImageInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2SetTimestepsStep"), + ("denoise.apply_strength", "Krea2ApplyStrengthStep"), + ("denoise.prepare_image_latents", "Krea2PrepareImageLatentsStep"), + ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), + ("denoise.denoise", "Krea2DenoiseStep"), + ("decode", "Krea2DecodeStep"), + ], + "inpainting": [ + ("text_encoder", "Krea2TextEncoderStep"), + ("vae_encoder.preprocess", "Krea2InpaintProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2VaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TextInputsStep"), + ("denoise.input.image_inputs", "Krea2ImageInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2SetTimestepsStep"), + ("denoise.apply_strength", "Krea2ApplyStrengthStep"), + ("denoise.prepare_inpaint_latents.add_noise", "Krea2PrepareImageLatentsStep"), + ("denoise.prepare_inpaint_latents.prepare_mask", "Krea2PrepareMaskLatentsStep"), + ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), + ("denoise.denoise", "Krea2InpaintDenoiseStep"), + ("decode", "Krea2InpaintDecodeStep"), + ], + "reference": [ + ("text_encoder", "Krea2ReferenceTextEncoderStep"), + ("vae_encoder.preprocess", "Krea2ReferenceProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2ReferenceVaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TextInputsStep"), + ("denoise.input.reference_inputs", "Krea2ReferenceInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2SetTimestepsStep"), + ("denoise.prepare_position_ids", "Krea2PrepareReferencePositionIdsStep"), + ("denoise.denoise", "Krea2ReferenceDenoiseStep"), + ("decode", "Krea2DecodeStep"), + ], } @@ -37,8 +81,8 @@ class TestKrea2ModularPipelineFast(ModularPipelineTesterMixin): pipeline_blocks_class = Krea2AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-modular-pipe" - params = frozenset(["prompt", "height", "width"]) - batch_params = frozenset(["prompt"]) + params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image", "reference_image_2"]) + batch_params = frozenset(["prompt", "image", "mask_image", "reference_image", "reference_image_2"]) expected_workflow_blocks = KREA2_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -56,3 +100,34 @@ def get_dummy_inputs(self, seed=0): def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) + + def test_image2image(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["image"] = PIL.Image.new("RGB", (32, 32), "white") + output = pipe(**inputs, strength=0.8, output="images") + assert output.shape == (1, 3, 32, 32) + + def test_inpainting(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["image"] = PIL.Image.new("RGB", (32, 32), "white") + inputs["mask_image"] = PIL.Image.new("L", (32, 32), "black") + output_low_strength = pipe(**inputs, strength=0.5, output="images") + inputs["generator"] = self.get_generator(0) + output_full_strength = pipe(**inputs, strength=1.0, output="images") + assert output_low_strength.shape == (1, 3, 32, 32) + assert (output_low_strength - output_full_strength).abs().max() < 1e-6 + + def test_reference_image(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["reference_image"] = PIL.Image.new("RGB", (32, 32), "white") + inputs["reference_image_encoder_resolution"] = 32 + single_reference_output = pipe(**inputs, reference_attention_scale=2.0, output="images") + + inputs["generator"] = self.get_generator(0) + inputs["reference_image_2"] = PIL.Image.new("RGB", (32, 32), "black") + two_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0], output="images") + assert single_reference_output.shape == (1, 3, 32, 32) + assert not torch.allclose(single_reference_output, two_reference_output) diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py index 9143e67d003a..a772fd57cf76 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py @@ -14,6 +14,9 @@ # limitations under the License. +import PIL +import torch + from diffusers.modular_pipelines import Krea2TurboAutoBlocks, Krea2TurboModularPipeline from ..test_modular_pipelines_common import ModularPipelineTesterMixin @@ -29,6 +32,47 @@ ("denoise.denoise", "Krea2TurboDenoiseStep"), ("decode", "Krea2DecodeStep"), ], + "image2image": [ + ("text_encoder", "Krea2TurboTextEncoderStep"), + ("vae_encoder.preprocess", "Krea2ProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2VaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TurboTextInputsStep"), + ("denoise.input.image_inputs", "Krea2ImageInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2TurboSetTimestepsStep"), + ("denoise.apply_strength", "Krea2ApplyStrengthStep"), + ("denoise.prepare_image_latents", "Krea2PrepareImageLatentsStep"), + ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), + ("denoise.denoise", "Krea2TurboDenoiseStep"), + ("decode", "Krea2DecodeStep"), + ], + "inpainting": [ + ("text_encoder", "Krea2TurboTextEncoderStep"), + ("vae_encoder.preprocess", "Krea2InpaintProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2VaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TurboTextInputsStep"), + ("denoise.input.image_inputs", "Krea2ImageInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2TurboSetTimestepsStep"), + ("denoise.apply_strength", "Krea2ApplyStrengthStep"), + ("denoise.prepare_inpaint_latents.add_noise", "Krea2PrepareImageLatentsStep"), + ("denoise.prepare_inpaint_latents.prepare_mask", "Krea2PrepareMaskLatentsStep"), + ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), + ("denoise.denoise", "Krea2TurboInpaintDenoiseStep"), + ("decode", "Krea2InpaintDecodeStep"), + ], + "reference": [ + ("text_encoder", "Krea2TurboReferenceTextEncoderStep"), + ("vae_encoder.preprocess", "Krea2ReferenceProcessImagesInputStep"), + ("vae_encoder.encode", "Krea2ReferenceVaeEncoderStep"), + ("denoise.input.text_inputs", "Krea2TurboTextInputsStep"), + ("denoise.input.reference_inputs", "Krea2ReferenceInputsStep"), + ("denoise.prepare_latents", "Krea2PrepareLatentsStep"), + ("denoise.set_timesteps", "Krea2TurboSetTimestepsStep"), + ("denoise.prepare_position_ids", "Krea2PrepareReferencePositionIdsStep"), + ("denoise.denoise", "Krea2TurboReferenceDenoiseStep"), + ("decode", "Krea2DecodeStep"), + ], } @@ -37,8 +81,8 @@ class TestKrea2TurboModularPipelineFast(ModularPipelineTesterMixin): pipeline_blocks_class = Krea2TurboAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-turbo-modular-pipe" - params = frozenset(["prompt", "height", "width"]) - batch_params = frozenset(["prompt"]) + params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image", "reference_image_2"]) + batch_params = frozenset(["prompt", "image", "mask_image", "reference_image", "reference_image_2"]) expected_workflow_blocks = KREA2_TURBO_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -56,3 +100,34 @@ def get_dummy_inputs(self, seed=0): def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) + + def test_image2image(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["image"] = PIL.Image.new("RGB", (32, 32), "white") + output = pipe(**inputs, strength=0.8, output="images") + assert output.shape == (1, 3, 32, 32) + + def test_inpainting(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["image"] = PIL.Image.new("RGB", (32, 32), "white") + inputs["mask_image"] = PIL.Image.new("L", (32, 32), "black") + output_low_strength = pipe(**inputs, strength=0.5, output="images") + inputs["generator"] = self.get_generator(0) + output_full_strength = pipe(**inputs, strength=1.0, output="images") + assert output_low_strength.shape == (1, 3, 32, 32) + assert (output_low_strength - output_full_strength).abs().max() < 1e-6 + + def test_reference_image(self): + pipe = self.get_pipeline().to("cpu") + inputs = self.get_dummy_inputs() + inputs["reference_image"] = PIL.Image.new("RGB", (32, 32), "white") + inputs["reference_image_encoder_resolution"] = 32 + single_reference_output = pipe(**inputs, reference_attention_scale=2.0, output="images") + + inputs["generator"] = self.get_generator(0) + inputs["reference_image_2"] = PIL.Image.new("RGB", (32, 32), "black") + two_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0], output="images") + assert single_reference_output.shape == (1, 3, 32, 32) + assert not torch.allclose(single_reference_output, two_reference_output) From e67bd2caf337868c70259ad6e4edda5101af2864 Mon Sep 17 00:00:00 2001 From: lucasruan1618 Date: Mon, 3 Aug 2026 16:19:31 +0000 Subject: [PATCH 2/2] Generalize Krea2 reference conditioning to multiple images --- docs/source/en/api/pipelines/krea2.md | 15 ++-- .../modular_pipelines/krea2/encoders.py | 86 +++++++------------ .../krea2/modular_blocks_krea2.py | 16 +--- .../krea2/modular_blocks_krea2_turbo.py | 8 +- .../test_models_transformer_krea2.py | 8 +- .../krea2/test_modular_pipeline_krea2.py | 25 ++++-- .../test_modular_pipeline_krea2_turbo.py | 25 ++++-- 7 files changed, 87 insertions(+), 96 deletions(-) diff --git a/docs/source/en/api/pipelines/krea2.md b/docs/source/en/api/pipelines/krea2.md index 2142cdc88124..96b4fc582f0d 100644 --- a/docs/source/en/api/pipelines/krea2.md +++ b/docs/source/en/api/pipelines/krea2.md @@ -190,8 +190,7 @@ subject_image = load_image( ) image = pipe( prompt="place the wizard cat from the second image sitting on the bench beside the dog from the first image", - reference_image=scene_image, - reference_image_2=subject_image, + reference_image=[scene_image, subject_image], height=1024, width=1024, reference_image_encoder_resolution=768, @@ -202,12 +201,12 @@ image = pipe( image.save("krea2_reference.png") ``` -For two-reference generation, `reference_image` is the scene and `reference_image_2` is the subject, matching the -adapter's training order. `reference_image_encoder_resolution` controls the maximum reference-image side length passed -to Qwen3-VL. `reference_attention_scale` accepts either one value for all references or one value per reference; the -example leaves scene attention unchanged and boosts subject fidelity. The adapter's recommended LoRA scale is `1.0`. -References are resized to the requested output dimensions before VAE encoding, so use similar aspect ratios to avoid -distortion. +`reference_image` accepts one image or an ordered list of any length. The example passes the scene first and the subject +second to match the adapter's training order. The same reference set is shared by every prompt in a prompt batch. +`reference_image_encoder_resolution` controls the maximum reference-image side length passed to Qwen3-VL. +`reference_attention_scale` accepts either one value for all references or one value per reference; the example leaves +scene attention unchanged and boosts subject fidelity. The adapter's recommended LoRA scale is `1.0`. References are +resized to the requested output dimensions before VAE encoding, so use similar aspect ratios to avoid distortion. We additionally provide an example for using Krea2 Turbo. The distilled checkpoint maps to its own set of blocks ([`Krea2TurboAutoBlocks`]): it runs guidance-free (no `guider`), takes no negative prompt, and samples in a few steps. diff --git a/src/diffusers/modular_pipelines/krea2/encoders.py b/src/diffusers/modular_pipelines/krea2/encoders.py index 9d9af2729079..cd45e5e3debb 100644 --- a/src/diffusers/modular_pipelines/krea2/encoders.py +++ b/src/diffusers/modular_pipelines/krea2/encoders.py @@ -318,7 +318,7 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi # auto_docstring class Krea2ReferenceTextEncoderStep(ModularPipelineBlocks): """ - Encode prompts together with a reference image through Qwen3-VL for reference-conditioned Krea 2 generation. + Encode prompts with one or more ordered reference images through Qwen3-VL for Krea 2 generation. Components: text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor @@ -331,9 +331,7 @@ class Krea2ReferenceTextEncoderStep(ModularPipelineBlocks): negative_prompt (`str`, *optional*): The negative prompt(s) for CFG. reference_image (`Image | list`): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. @@ -352,7 +350,7 @@ class Krea2ReferenceTextEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: - return "Encode prompts together with a reference image through Qwen3-VL for reference-conditioned Krea 2 generation." + return "Encode prompts with one or more ordered reference images through Qwen3-VL for Krea 2 generation." @property def expected_components(self) -> list[ComponentSpec]: @@ -382,12 +380,7 @@ def inputs(self) -> list[InputParam]: name="reference_image", type_hint=PIL.Image.Image | list[PIL.Image.Image], required=True, - description="First reference image(s), or scene reference for two-reference generation.", - ), - InputParam( - name="reference_image_2", - type_hint=PIL.Image.Image | list[PIL.Image.Image], - description="Optional second reference image(s), used as the subject reference.", + description="A reference image or ordered list of reference images shared by all prompts in the batch.", ), InputParam( name="reference_image_encoder_resolution", @@ -406,25 +399,17 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam.template("negative_prompt_embeds_mask"), ] - def _encode_prompt(self, components, prompts, reference_images, reference_images_2, encoder_resolution, device): - references_by_input = [] - for name, images in (("reference_image", reference_images), ("reference_image_2", reference_images_2)): - if images is None: - continue - if isinstance(images, PIL.Image.Image): - images = [images] - if len(images) == 1 and len(prompts) > 1: - images = images * len(prompts) - if len(images) != len(prompts): - raise ValueError( - f"`{name}` must contain one image or one image per prompt, but got {len(images)} images for " - f"{len(prompts)} prompts." - ) - references_by_input.append(images) + def _encode_prompt(self, components, prompts, reference_images, encoder_resolution, device): + if isinstance(reference_images, PIL.Image.Image): + reference_images = [reference_images] + if not isinstance(reference_images, list) or not reference_images: + raise ValueError("`reference_image` must be an image or a non-empty list of images.") + if not all(isinstance(image, PIL.Image.Image) for image in reference_images): + raise ValueError("Every item in `reference_image` must be a PIL image.") processed_images = [] - for prompt_images in zip(*references_by_input): - for image in prompt_images: + for _ in prompts: + for image in reference_images: image = image.convert("RGB") if encoder_resolution and max(image.size) > encoder_resolution: scale = encoder_resolution / max(image.size) @@ -439,7 +424,7 @@ def _encode_prompt(self, components, prompts, reference_images, reference_images image_token_counts = ( image_inputs.image_grid_thw.prod(dim=1) // components.reference_image_processor.merge_size**2 ).tolist() - num_references = len(references_by_input) + num_references = len(reference_images) vision_block = "<|vision_start|><|image_pad|><|vision_end|>" texts = [] for prompt_index, prompt in enumerate(prompts): @@ -476,7 +461,6 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi components, prompts, block_state.reference_image, - block_state.reference_image_2, block_state.reference_image_encoder_resolution, device, ) @@ -493,7 +477,6 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi components, negative_prompts, block_state.reference_image, - block_state.reference_image_2, block_state.reference_image_encoder_resolution, device, ) @@ -523,7 +506,7 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi # auto_docstring class Krea2TurboReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): """ - Encode prompts with a reference image for reference-conditioned Krea 2 Turbo generation. + Encode prompts with one or more ordered reference images for Krea 2 Turbo generation. Components: text_encoder (`Qwen3VLModel`): The Qwen3-VL text encoder. reference_image_processor @@ -534,9 +517,7 @@ class Krea2TurboReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): prompt (`str`): The prompt or prompts to guide image generation. reference_image (`Image | list`): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. @@ -549,7 +530,7 @@ class Krea2TurboReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep): @property def description(self) -> str: - return "Encode prompts with a reference image for reference-conditioned Krea 2 Turbo generation." + return "Encode prompts with one or more ordered reference images for Krea 2 Turbo generation." @property def expected_components(self) -> list[ComponentSpec]: @@ -572,12 +553,7 @@ def inputs(self) -> list[InputParam]: name="reference_image", type_hint=PIL.Image.Image | list[PIL.Image.Image], required=True, - description="First reference image(s), or scene reference for two-reference generation.", - ), - InputParam( - name="reference_image_2", - type_hint=PIL.Image.Image | list[PIL.Image.Image], - description="Optional second reference image(s), used as the subject reference.", + description="A reference image or ordered list of reference images shared by all prompts in the batch.", ), InputParam( name="reference_image_encoder_resolution", @@ -599,7 +575,6 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi components, prompts, block_state.reference_image, - block_state.reference_image_2, block_state.reference_image_encoder_resolution, components._execution_device, ) @@ -841,16 +816,14 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi # auto_docstring class Krea2ReferenceProcessImagesInputStep(ModularPipelineBlocks): """ - Preprocess a reference image at the target output resolution for VAE encoding. + Preprocess one or more ordered reference images at the target output resolution for VAE encoding. Components: image_processor (`VaeImageProcessor`) Inputs: reference_image (`Image | list`): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. width (`int`, *optional*, defaults to 1024): @@ -865,7 +838,7 @@ class Krea2ReferenceProcessImagesInputStep(ModularPipelineBlocks): @property def description(self) -> str: - return "Preprocess a reference image at the target output resolution for VAE encoding." + return "Preprocess one or more ordered reference images at the target output resolution for VAE encoding." @property def expected_components(self) -> list[ComponentSpec]: @@ -885,12 +858,7 @@ def inputs(self) -> list[InputParam]: name="reference_image", type_hint=PIL.Image.Image | list[PIL.Image.Image], required=True, - description="First reference image(s), or scene reference for two-reference generation.", - ), - InputParam( - name="reference_image_2", - type_hint=PIL.Image.Image | list[PIL.Image.Image], - description="Optional second reference image(s), used as the subject reference.", + description="A reference image or ordered list of reference images shared by all prompts in the batch.", ), InputParam.template("height", default=1024), InputParam.template("width", default=1024), @@ -912,9 +880,13 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi multiple = components.image_processor.config.vae_scale_factor if block_state.height % multiple != 0 or block_state.width % multiple != 0: raise ValueError(f"`height` and `width` must be divisible by {multiple} for reference conditioning.") - reference_images = [block_state.reference_image] - if block_state.reference_image_2 is not None: - reference_images.append(block_state.reference_image_2) + reference_images = block_state.reference_image + if isinstance(reference_images, PIL.Image.Image): + reference_images = [reference_images] + if not isinstance(reference_images, list) or not reference_images: + raise ValueError("`reference_image` must be an image or a non-empty list of images.") + if not all(isinstance(image, PIL.Image.Image) for image in reference_images): + raise ValueError("Every item in `reference_image` must be a PIL image.") block_state.processed_reference_images = [ components.image_processor.preprocess(image=image, height=block_state.height, width=block_state.width) for image in reference_images diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py index 0e067b40b02d..44e40c036a02 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py +++ b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py @@ -60,9 +60,7 @@ class Krea2AutoTextEncoderStep(AutoPipelineBlocks): negative_prompt (`str`, *optional*): The negative prompt(s) for CFG. reference_image (`Image | list`, *optional*): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): @@ -242,9 +240,7 @@ class Krea2ReferenceVaeEncoderBlocks(SequentialPipelineBlocks): Inputs: reference_image (`Image | list`): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. height (`int`, *optional*, defaults to 1024): The height in pixels of the generated image. width (`int`, *optional*, defaults to 1024): @@ -277,9 +273,7 @@ class Krea2AutoVaeEncoderStep(AutoPipelineBlocks): Inputs: reference_image (`Image | list`, *optional*): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. height (`int`, *optional*, defaults to 1024 or None, depending on the workflow): The height in pixels of the generated image. width (`int`, *optional*, defaults to 1024 or None, depending on the workflow): @@ -812,9 +806,7 @@ class Krea2AutoBlocks(SequentialPipelineBlocks): negative_prompt (`str`, *optional*): The negative prompt(s) for CFG. reference_image (`Image | list`, *optional*): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py index 3568b1430a18..525c2a98cbe9 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py +++ b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py @@ -53,9 +53,7 @@ class Krea2TurboAutoTextEncoderStep(AutoPipelineBlocks): prompt (`str`): The prompt or prompts to guide image generation. reference_image (`Image | list`, *optional*): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): @@ -526,9 +524,7 @@ class Krea2TurboAutoBlocks(SequentialPipelineBlocks): prompt (`str`): The prompt or prompts to guide image generation. reference_image (`Image | list`, *optional*): - First reference image(s), or scene reference for two-reference generation. - reference_image_2 (`Image | list`, *optional*): - Optional second reference image(s), used as the subject reference. + A reference image or ordered list of reference images shared by all prompts in the batch. reference_image_encoder_resolution (`int`, *optional*, defaults to 768): Maximum reference-image side length used by the Qwen3-VL encoder. Use 0 for native resolution. max_sequence_length (`int`, *optional*, defaults to 512): diff --git a/tests/models/transformers/test_models_transformer_krea2.py b/tests/models/transformers/test_models_transformer_krea2.py index 6b5cea8ff27f..4a2cb06b173d 100644 --- a/tests/models/transformers/test_models_transformer_krea2.py +++ b/tests/models/transformers/test_models_transformer_krea2.py @@ -143,22 +143,24 @@ def test_reference_hidden_states(self): device=torch_device, dtype=self.torch_dtype, ) - for _ in range(2) + for _ in range(3) ] position_ids = torch.cat( [ inputs["position_ids"][:text_seq_len], inputs["position_ids"][text_seq_len:].clone(), inputs["position_ids"][text_seq_len:].clone(), + inputs["position_ids"][text_seq_len:].clone(), inputs["position_ids"][text_seq_len:], ] ) position_ids[text_seq_len : text_seq_len + image_seq_len, 0] = 1 position_ids[text_seq_len + image_seq_len : text_seq_len + 2 * image_seq_len, 0] = 2 + position_ids[text_seq_len + 2 * image_seq_len : text_seq_len + 3 * image_seq_len, 0] = 3 inputs["position_ids"] = position_ids - output = model(**inputs, reference_attention_scale=[1.0, 1.0]).sample - boosted_output = model(**inputs, reference_attention_scale=[1.0, 2.0]).sample + output = model(**inputs, reference_attention_scale=[1.0, 1.0, 1.0]).sample + boosted_output = model(**inputs, reference_attention_scale=[1.0, 2.0, 0.5]).sample assert output.shape == inputs["hidden_states"].shape assert not torch.allclose(output, boosted_output) diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index b8cd6c45ea08..fd6462d8fdf4 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -81,8 +81,8 @@ class TestKrea2ModularPipelineFast(ModularPipelineTesterMixin): pipeline_blocks_class = Krea2AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-modular-pipe" - params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image", "reference_image_2"]) - batch_params = frozenset(["prompt", "image", "mask_image", "reference_image", "reference_image_2"]) + params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image"]) + batch_params = frozenset(["prompt", "image", "mask_image"]) expected_workflow_blocks = KREA2_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -127,7 +127,22 @@ def test_reference_image(self): single_reference_output = pipe(**inputs, reference_attention_scale=2.0, output="images") inputs["generator"] = self.get_generator(0) - inputs["reference_image_2"] = PIL.Image.new("RGB", (32, 32), "black") - two_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0], output="images") + inputs["reference_image"] = [PIL.Image.new("RGB", (32, 32), "white")] + single_reference_list_output = pipe(**inputs, reference_attention_scale=[2.0], output="images") + + inputs["generator"] = self.get_generator(0) + inputs["reference_image"] = [ + PIL.Image.new("RGB", (32, 32), "white"), + PIL.Image.new("RGB", (32, 32), "black"), + PIL.Image.new("RGB", (32, 32), "gray"), + ] + multi_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0, 0.5], output="images") + + inputs["prompt"] = [inputs["prompt"], inputs["prompt"]] + inputs["generator"] = [self.get_generator(0), self.get_generator(1)] + batched_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0, 0.5], output="images") assert single_reference_output.shape == (1, 3, 32, 32) - assert not torch.allclose(single_reference_output, two_reference_output) + assert torch.allclose(single_reference_output, single_reference_list_output) + assert not torch.allclose(single_reference_output, multi_reference_output) + assert batched_output.shape == (2, 3, 32, 32) + assert (batched_output[:1] - multi_reference_output).abs().max() < 5e-3 diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py index a772fd57cf76..1df448be6c30 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py @@ -81,8 +81,8 @@ class TestKrea2TurboModularPipelineFast(ModularPipelineTesterMixin): pipeline_blocks_class = Krea2TurboAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-turbo-modular-pipe" - params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image", "reference_image_2"]) - batch_params = frozenset(["prompt", "image", "mask_image", "reference_image", "reference_image_2"]) + params = frozenset(["prompt", "height", "width", "image", "mask_image", "reference_image"]) + batch_params = frozenset(["prompt", "image", "mask_image"]) expected_workflow_blocks = KREA2_TURBO_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -127,7 +127,22 @@ def test_reference_image(self): single_reference_output = pipe(**inputs, reference_attention_scale=2.0, output="images") inputs["generator"] = self.get_generator(0) - inputs["reference_image_2"] = PIL.Image.new("RGB", (32, 32), "black") - two_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0], output="images") + inputs["reference_image"] = [PIL.Image.new("RGB", (32, 32), "white")] + single_reference_list_output = pipe(**inputs, reference_attention_scale=[2.0], output="images") + + inputs["generator"] = self.get_generator(0) + inputs["reference_image"] = [ + PIL.Image.new("RGB", (32, 32), "white"), + PIL.Image.new("RGB", (32, 32), "black"), + PIL.Image.new("RGB", (32, 32), "gray"), + ] + multi_reference_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0, 0.5], output="images") + + inputs["prompt"] = [inputs["prompt"], inputs["prompt"]] + inputs["generator"] = [self.get_generator(0), self.get_generator(1)] + batched_output = pipe(**inputs, reference_attention_scale=[1.0, 2.0, 0.5], output="images") assert single_reference_output.shape == (1, 3, 32, 32) - assert not torch.allclose(single_reference_output, two_reference_output) + assert torch.allclose(single_reference_output, single_reference_list_output) + assert not torch.allclose(single_reference_output, multi_reference_output) + assert batched_output.shape == (2, 3, 32, 32) + assert (batched_output[:1] - multi_reference_output).abs().max() < 5e-3