diff --git a/docs/source/en/api/pipelines/krea2.md b/docs/source/en/api/pipelines/krea2.md index b6e6fd5998c3..96b4fc582f0d 100644 --- a/docs/source/en/api/pipelines/krea2.md +++ b/docs/source/en/api/pipelines/krea2.md @@ -106,6 +106,108 @@ 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, 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") +``` + +`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. `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..cd45e5e3debb 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,650 @@ 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 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 + (`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`): + 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. + + 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 with one or more ordered reference images through Qwen3-VL for 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="A reference image or ordered list of reference images shared by all prompts in the batch.", + ), + 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, 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 _ 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) + 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(reference_images) + 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_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_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 one or more ordered reference images for 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`): + 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. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + prompt_embeds_mask (`Tensor`): + The encoder attention mask. + """ + + @property + def description(self) -> str: + return "Encode prompts with one or more ordered reference images for 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="A reference image or ordered list of reference images shared by all prompts in the batch.", + ), + 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_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 one or more ordered reference images at the target output resolution for VAE encoding. + + Components: + image_processor (`VaeImageProcessor`) + + Inputs: + reference_image (`Image | list`): + 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): + 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 one or more ordered reference images 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="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), + ] + + @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 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 + ] + 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..44e40c036a02 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py +++ b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py @@ -14,22 +14,79 @@ 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*): + 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): + 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 +158,695 @@ 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`): + 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): + 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*): + 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): + 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*): + 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): 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 +855,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..525c2a98cbe9 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,68 @@ 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*): + 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): + 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 +143,428 @@ 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*): + 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): 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 +573,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..4a2cb06b173d 100644 --- a/tests/models/transformers/test_models_transformer_krea2.py +++ b/tests/models/transformers/test_models_transformer_krea2.py @@ -131,6 +131,39 @@ 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(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, 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) + 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..fd6462d8fdf4 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"]) + batch_params = frozenset(["prompt", "image", "mask_image"]) expected_workflow_blocks = KREA2_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -56,3 +100,49 @@ 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"] = [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 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 9143e67d003a..1df448be6c30 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"]) + batch_params = frozenset(["prompt", "image", "mask_image"]) expected_workflow_blocks = KREA2_TURBO_WORKFLOWS def get_dummy_inputs(self, seed=0): @@ -56,3 +100,49 @@ 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"] = [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 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