[Modular] Krea2 Img2Img, Inpaint and Reference Images modular pipeline

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

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: