krea-2 reference blocks

.gitignore

@@ -0,0 +1,2 @@
+__pycache__/
+.ruff_cache/

README.md

@@ -1,3 +1,13 @@
 ---
+library_name: diffusers
+pipeline_tag: text-to-image
+base_model: krea/Krea-2-Raw
+tags:
+  - text-to-image
+  - image-to-image
+  - reference-image
+  - modular-diffusers
+  - diffusion
+  - krea
 license: apache-2.0
 ---

__init__.py

@@ -0,0 +1,60 @@
+from .krea2_reference import (
+    DEFAULT_GROUNDING_PX,
+    REFERENCE_MASK_MODE_DEEMPHASIZE,
+    REFERENCE_MASK_MODE_EXCLUDE,
+    REFERENCE_MASK_MODE_EXCLUDE_BLANK,
+    REFERENCE_MASK_MODES,
+    REFERENCE_MODE_APPEND,
+    REFERENCE_MODE_OFF,
+    REFERENCE_MODE_PREPEND,
+    REFERENCE_MODES,
+    Krea2AppendReferenceImagesStep,
+    Krea2AppendReferenceTextEncoderStep,
+    Krea2AppendReferenceWorkflow,
+    Krea2PrependReferenceImagesStep,
+    Krea2PrependReferenceTextEncoderStep,
+    Krea2PrependReferenceWorkflow,
+    Krea2ReferenceAutoBlocks,
+    Krea2ReferenceLatentsStep,
+    Krea2ReferenceTextInputsStep,
+    Krea2TurboAppendReferenceWorkflow,
+    Krea2TurboPrependReferenceWorkflow,
+    Krea2TurboReferenceAutoBlocks,
+    Krea2TurboVisionReferenceWorkflow,
+    Krea2VisionReferenceImagesStep,
+    Krea2VisionReferenceTextEncoderStep,
+    Krea2VisionReferenceWorkflow,
+)
+
+
+__all__ = [
+    # Top level: `reference_mode` selects the workflow.
+    "Krea2ReferenceAutoBlocks",
+    "Krea2TurboReferenceAutoBlocks",
+    # Per-mode workflows.
+    "Krea2VisionReferenceWorkflow",
+    "Krea2AppendReferenceWorkflow",
+    "Krea2PrependReferenceWorkflow",
+    "Krea2TurboVisionReferenceWorkflow",
+    "Krea2TurboAppendReferenceWorkflow",
+    "Krea2TurboPrependReferenceWorkflow",
+    # Leaf blocks.
+    "Krea2VisionReferenceImagesStep",
+    "Krea2AppendReferenceImagesStep",
+    "Krea2PrependReferenceImagesStep",
+    "Krea2VisionReferenceTextEncoderStep",
+    "Krea2AppendReferenceTextEncoderStep",
+    "Krea2PrependReferenceTextEncoderStep",
+    "Krea2ReferenceTextInputsStep",
+    "Krea2ReferenceLatentsStep",
+    # Mode constants.
+    "DEFAULT_GROUNDING_PX",
+    "REFERENCE_MODES",
+    "REFERENCE_MODE_APPEND",
+    "REFERENCE_MODE_OFF",
+    "REFERENCE_MODE_PREPEND",
+    "REFERENCE_MASK_MODES",
+    "REFERENCE_MASK_MODE_DEEMPHASIZE",
+    "REFERENCE_MASK_MODE_EXCLUDE",
+    "REFERENCE_MASK_MODE_EXCLUDE_BLANK",
+]

krea2_reference.py

@@ -0,0 +1,1790 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Custom modular-diffusers blocks adding reference-image conditioning to Krea 2.
+
+Krea 2's text encoder is the full Qwen3-VL VLM, vision tower included, so reference images can condition
+generation through the encoder's vision path. Two community edit LoRAs also feed the references as *clean VAE
+latents* inside the transformer sequence. All three paths live here, selected by `reference_mode`:
+
+* `"off"` (default) -- vision path only. Works on the stock base checkpoint, no LoRA.
+* `"append"` -- vision (coarse view) + clean VAE reference tokens **after** the noisy target, modulated at flow
+  time `t=0`. The Ostris AI-Toolkit reference/edit LoRA recipe.
+* `"prepend"` -- vision (grounded view) + the clean VAE source **before** the noisy target under the plain
+  uniform timestep modulation, distinguished only by its RoPE frame axis. The Identity-Edit recipe
+  (ai-toolkit `predict_velocity_edit`); its unconditional branch is grounded on the same image.
+
+With `reference_mode="off"` and no images the blocks reduce to stock text-to-image. No mode patches the
+`Krea2Transformer2DModel` source.
+"""
+
+import math
+from typing import Any
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn.functional as F
+from transformers import Qwen3VLProcessor
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.guiders import ClassifierFreeGuidance
+from diffusers.models import AutoencoderKLQwenImage
+from diffusers.models.transformers.transformer_krea2 import Krea2Transformer2DModel
+from diffusers.modular_pipelines.krea2.before_denoise import (
+    Krea2PrepareLatentsStep,
+    Krea2PreparePositionIdsStep,
+    Krea2SetTimestepsStep,
+    Krea2TurboSetTimestepsStep,
+)
+from diffusers.modular_pipelines.krea2.decoders import Krea2DecodeStep
+from diffusers.modular_pipelines.krea2.denoise import (
+    Krea2DenoiseLoopWrapper,
+    Krea2LoopAfterDenoiser,
+    Krea2LoopBeforeDenoiser,
+)
+from diffusers.modular_pipelines.krea2.encoders import (
+    _PROMPT_TEMPLATE_ENCODE_PREFIX,
+    _PROMPT_TEMPLATE_ENCODE_START_IDX,
+    _PROMPT_TEMPLATE_ENCODE_SUFFIX,
+    KREA2_TEXT_ENCODER_SELECT_LAYERS,
+    Krea2TextEncoderStep,
+)
+from diffusers.modular_pipelines.krea2.modular_pipeline import Krea2ModularPipeline
+from diffusers.modular_pipelines.modular_pipeline import (
+    BlockState,
+    ConditionalPipelineBlocks,
+    ModularPipelineBlocks,
+    PipelineState,
+    SequentialPipelineBlocks,
+)
+from diffusers.modular_pipelines.modular_pipeline_utils import (
+    ComponentSpec,
+    InputParam,
+    InsertableDict,
+    OutputParam,
+)
+from diffusers.utils import logging
+
+from .krea2_vision_attention import build_patch_keeps, vision_key_masks
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+REFERENCE_MODE_OFF = "off"
+REFERENCE_MODE_APPEND = "append"
+REFERENCE_MODE_PREPEND = "prepend"
+REFERENCE_MODES = (REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND)
+
+# What a painted mask does to the Qwen3-VL vision path. Zeroing output tokens alone cannot remove content
+# from a reference (see `krea2_vision_attention`), so `deemphasize` attenuates rather than excludes.
+REFERENCE_MASK_MODE_EXCLUDE_BLANK = "exclude_blank"
+REFERENCE_MASK_MODE_EXCLUDE = "exclude"
+REFERENCE_MASK_MODE_DEEMPHASIZE = "deemphasize"
+REFERENCE_MASK_MODES = (
+    REFERENCE_MASK_MODE_EXCLUDE_BLANK,
+    REFERENCE_MASK_MODE_EXCLUDE,
+    REFERENCE_MASK_MODE_DEEMPHASIZE,
+)
+
+# Qwen3-VL vision-path budgets. `off` matches the reference ComfyUI encoder's default; `append` matches the
+# AI-Toolkit recipe. `prepend` sizes its view from `grounding_px` instead.
+VISION_MAX_PIXELS = 1024 * 1024
+VISION_EDIT_MAX_PIXELS = 384 * 384
+
+REFERENCE_LATENTS_MAX_PIXELS = 1024 * 1024
+
+# Longest-side cap (px) for the `prepend` grounding view; 0 means native. The Identity-Edit LoRA trained on
+# 384-768px jitter.
+DEFAULT_GROUNDING_PX = 768
+
+
+def _coerce_reference_mode(value: Any) -> str:
+    """Normalize a reference mode, falling back to `"off"` for anything unrecognized."""
+    if isinstance(value, str) and value in REFERENCE_MODES:
+        return value
+    if value is not None and value != REFERENCE_MODE_OFF:
+        logger.warning(f"Unknown `reference_mode` {value!r}; falling back to {REFERENCE_MODE_OFF!r}.")
+    return REFERENCE_MODE_OFF
+
+
+def _coerce_reference_mask_mode(value: Any) -> str:
+    """Normalize a reference-mask mode, falling back to `"exclude_blank"` for anything unrecognized."""
+    if isinstance(value, str) and value in REFERENCE_MASK_MODES:
+        return value
+    if value is not None:
+        logger.warning(
+            f"Unknown `reference_mask_mode` {value!r}; falling back to {REFERENCE_MASK_MODE_EXCLUDE_BLANK!r}."
+        )
+    return REFERENCE_MASK_MODE_EXCLUDE_BLANK
+
+
+def _as_list(value: Any) -> list:
+    """Wrap a bare value into a single-element list; `None` becomes an empty list."""
+    if value is None:
+        return []
+    if isinstance(value, (list, tuple)):
+        return list(value)
+    return [value]
+
+
+def _broadcast_per_image(value: Any, num_images: int, default: float) -> list[float]:
+    """Expand a scalar (or short list) of per-image strengths to one float per reference image."""
+    values = _as_list(value)
+    if not values:
+        return [default] * num_images
+    if len(values) == 1:
+        return [float(values[0])] * num_images
+    values = [float(v) for v in values[:num_images]]
+    return values + [default] * (num_images - len(values))
+
+
+def _to_rgb(image) -> PIL.Image.Image:
+    if not isinstance(image, PIL.Image.Image):
+        image = PIL.Image.fromarray(np.asarray(image))
+    return image.convert("RGB")
+
+
+def _cap_area(image, max_pixels: int) -> PIL.Image.Image:
+    """Aspect-preserving downscale so the image covers at most `max_pixels`. Never upscales."""
+    image = _to_rgb(image)
+    pixels = image.width * image.height
+    if pixels > max_pixels:
+        scale = math.sqrt(max_pixels / pixels)
+        image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))))
+    return image
+
+
+def _cap_longest_side(image, longest_side: int) -> PIL.Image.Image:
+    """Downscale so the longest side is at most `longest_side` px (0 = native). Never upscales."""
+    image = _to_rgb(image)
+    longest = max(image.width, image.height)
+    if longest_side and longest > longest_side:
+        scale = longest_side / longest
+        image = image.resize((max(1, round(image.width * scale)), max(1, round(image.height * scale))))
+    return image
+
+
+def _mask_to_keep_grid(mask, grid_height: int, grid_width: int):
+    """Rasterize a painted mask onto a `(grid_height, grid_width)` bool grid, `True` = keep.
+
+    Painted is the part to *use*. Shared by every consumer so the polarity means the same thing on the
+    pre-merge patch grid, the merged vision-token grid and the VAE latent patch grid. `RGBA` masks are read
+    from their alpha channel, anything else from luminance. Returns `None` when the mask is absent or paints
+    nothing -- an empty mask read literally would drop every token and silently disable the reference.
+    """
+    if mask is None or grid_height <= 0 or grid_width <= 0:
+        return None
+    if not isinstance(mask, PIL.Image.Image):
+        mask = PIL.Image.fromarray(np.asarray(mask))
+    painted = np.array(mask)[:, :, 3] if mask.mode == "RGBA" else np.array(mask.convert("L"))
+    painted = PIL.Image.fromarray(painted.astype(np.uint8)).resize((grid_width, grid_height), PIL.Image.BILINEAR)
+
+    keep = np.asarray(painted, dtype=np.float32) / 255.0 >= 0.5
+    if not keep.any():
+        logger.warning(
+            "A reference mask paints nothing and is being ignored. `RGBA` masks are read from their alpha "
+            "channel, so loading one through a helper that flattens to RGB (such as `load_image`) discards it."
+        )
+        return None
+    return keep
+
+
+def _mask_to_keep_vector(mask, grid_height: int, grid_width: int, temporal: int = 1):
+    """Row-major flattened `_mask_to_keep_grid`, tiled over `temporal` frames.
+
+    Row-major is correct at merged-token resolution; the pre-merge patch sequence is not row-major and uses
+    `krea2_vision_attention.patch_keep_vector` instead.
+    """
+    keep = _mask_to_keep_grid(mask, grid_height, grid_width)
+    if keep is None:
+        return None
+    flat = keep.astype(np.float32).reshape(-1)
+    if temporal > 1:
+        flat = np.tile(flat, temporal)
+    return torch.from_numpy(flat)
+
+
+def _build_keep_vectors(masks: list, image_grid_thw, merge_size: int) -> list:
+    """Build one keep vector per row of `image_grid_thw` (or `None` where the slot has no effective mask)."""
+    if image_grid_thw is None:
+        return []
+    keep_vectors = []
+    for i in range(int(image_grid_thw.shape[0])):
+        mask = masks[i] if i < len(masks) else None
+        t, h, w = (int(x) for x in image_grid_thw[i].tolist())
+        keep_vectors.append(_mask_to_keep_vector(mask, h // merge_size, w // merge_size, t))
+    return keep_vectors
+
+
+def _scaled_get_image_features(original_fn, style_strengths, subject_strengths, merge_size, keep_vectors=None):
+    """Wrap `Qwen3VLModel.get_image_features` to scale each reference's vision features per channel.
+
+    Qwen3-VL injects each image twice: as scattered `<|image_pad|>` embeddings (`pooler_output`, a per-image
+    list) and as deepstack features (per-layer tensors re-injected at deeper LLM layers). `subject` scales the
+    former, `style` the latter; 0 removes that channel, 1.0 leaves it at full effect.
+
+    `keep_vectors` (one entry per image, or `None`) additionally zeroes the vision tokens outside the painted
+    region in both channels.
+    """
+
+    def strength_for(values, index):
+        return float(values[index]) if index < len(values) else 1.0
+
+    def wrapper(pixel_values, image_grid_thw, **kwargs):
+        outputs = original_fn(pixel_values, image_grid_thw, **kwargs)
+
+        def keep_for(index, like):
+            if not keep_vectors or index >= len(keep_vectors) or keep_vectors[index] is None:
+                return None
+            return keep_vectors[index].to(device=like.device, dtype=like.dtype)
+
+        pooled = getattr(outputs, "pooler_output", None)
+        if isinstance(pooled, (list, tuple)):
+            scaled = []
+            for i, features in enumerate(pooled):
+                features = features * strength_for(subject_strengths, i)
+                keep = keep_for(i, features)
+                if keep is not None:
+                    features = features * keep.unsqueeze(-1)
+                scaled.append(features)
+            outputs.pooler_output = type(pooled)(scaled)
+
+        deepstack = getattr(outputs, "deepstack_features", None)
+        if deepstack:
+            reference = deepstack[0]
+            counts = [
+                int(image_grid_thw[i].prod().item()) // (merge_size * merge_size)
+                for i in range(image_grid_thw.shape[0])
+            ]
+            blocks = []
+            for i, count in enumerate(counts):
+                block = torch.full(
+                    (count,), strength_for(style_strengths, i), device=reference.device, dtype=reference.dtype
+                )
+                keep = keep_for(i, reference)
+                if keep is not None:
+                    block = block * keep
+                blocks.append(block)
+            per_token = torch.cat(blocks).unsqueeze(-1)
+            outputs.deepstack_features = type(deepstack)(layer * per_token for layer in deepstack)
+
+        return outputs
+
+    return wrapper
+
+
+def _build_reference_processor(text_encoder, tokenizer):
+    """Build a `Qwen3VLProcessor` for reference-image conditioning.
+
+    Krea 2 ships no image-processor config, so one is derived here. Qwen3-VL reuses the Qwen2-VL image
+    processor; the video processor is unused but the processor class requires a non-`None` instance. The
+    patch/merge/temporal sizes MUST come from the model's own vision config (Qwen3-VL uses `patch_size=16`, not
+    the Qwen2-VL default of 14) or the patch-embed reshape fails.
+    """
+    from transformers import Qwen3VLProcessor, Qwen3VLVideoProcessor
+    from transformers.models.qwen2_vl.image_processing_qwen2_vl_fast import Qwen2VLImageProcessorFast
+
+    vision_config = text_encoder.config.vision_config
+    image_processor = Qwen2VLImageProcessorFast(
+        patch_size=vision_config.patch_size,
+        merge_size=vision_config.spatial_merge_size,
+        temporal_patch_size=vision_config.temporal_patch_size,
+    )
+    return Qwen3VLProcessor(
+        image_processor=image_processor, tokenizer=tokenizer, video_processor=Qwen3VLVideoProcessor()
+    )
+
+
+def _pack_latents(latents: torch.Tensor, patch_size: int) -> torch.Tensor:
+    """Pack spatial latents `(B, C, H, W)` into Krea 2's sequence `(B, (H/p) * (W/p), C * p * p)`."""
+    batch_size, channels, height, width = latents.shape
+    latents = latents.view(batch_size, channels, height // patch_size, patch_size, width // patch_size, patch_size)
+    latents = latents.permute(0, 2, 4, 1, 3, 5)
+    return latents.reshape(
+        batch_size, (height // patch_size) * (width // patch_size), channels * patch_size * patch_size
+    )
+
+
+def _reference_position_ids(index: int, grid_height: int, grid_width: int) -> torch.Tensor:
+    """`(grid_h * grid_w, 3)` rotary coordinates for reference `index`: frame axis `index + 1`, own y/x grid."""
+    ids = torch.zeros(grid_height, grid_width, 3)
+    ids[..., 0] = index + 1
+    ids[..., 1] = torch.arange(grid_height)[:, None]
+    ids[..., 2] = torch.arange(grid_width)[None, :]
+    return ids.reshape(-1, 3)
+
+
+def _register_zero_time_reference_hooks(transformer, timestep, batch_size, split, total):
+    """Modulate the trailing `total - split` reference tokens at flow time `t=0` for the `append` mode.
+
+    A `forward_pre_hook` swaps each block's broadcast `(B, 1, 6 * hidden)` timestep embedding for a per-token
+    `(B, total, 6 * hidden)` one carrying the real timestep on the text and target rows and the `t=0` embedding
+    on the reference rows. The block's `temb.unflatten(-1, (6, -1)) + scale_shift_table` already broadcasts over
+    the sequence axis, so no source patching is needed.
+
+    Returns the hook handles; the caller must remove them once the forward pass is done.
+    """
+    dtype = transformer.dtype
+
+    def time_mod(t):
+        return transformer.time_mod_proj(F.gelu(transformer.time_embed(t, dtype=dtype), approximate="tanh"))
+
+    temb = time_mod(timestep)
+    temb_zero = time_mod(torch.zeros_like(timestep))
+    per_token = torch.cat([temb.expand(batch_size, split, -1), temb_zero.expand(batch_size, total - split, -1)], dim=1)
+
+    def pre_hook(module, args, kwargs):
+        if len(args) >= 2:
+            return (args[0], per_token, *args[2:]), kwargs
+        return args, {**kwargs, "temb": per_token}
+
+    return [block.register_forward_pre_hook(pre_hook, with_kwargs=True) for block in transformer.transformer_blocks]
+
+
+def _concat_reference_latents(latents, reference_latents, mode, dtype):
+    """Splice the clean reference tokens into the noisy latents in the order the mode dictates.
+
+    Returns `(latents, reference_seq_len)`; `reference_seq_len` is 0 when the mode carries no latent channel.
+    """
+    if mode == REFERENCE_MODE_OFF or reference_latents is None:
+        return latents, 0
+    reference_latents = reference_latents.to(device=latents.device, dtype=dtype)
+    if reference_latents.shape[0] != latents.shape[0]:
+        reference_latents = reference_latents.expand(latents.shape[0], -1, -1)
+    if mode == REFERENCE_MODE_PREPEND:
+        return torch.cat([reference_latents, latents], dim=1), reference_latents.shape[1]
+    return torch.cat([latents, reference_latents], dim=1), reference_latents.shape[1]
+
+
+def _slice_reference_rows(noise_pred, mode, reference_seq_len, target_seq_len):
+    """Drop the reference rows from a prediction, wherever the mode put them."""
+    if not reference_seq_len:
+        return noise_pred
+    if mode == REFERENCE_MODE_PREPEND:
+        return noise_pred[:, reference_seq_len:]
+    return noise_pred[:, :target_seq_len]
+
+
+# auto_docstring
+class Krea2ReferenceImagesStep(ModularPipelineBlocks):
+    """Base reference collector. Not used directly -- pick the subclass for the mode."""
+
+    model_name = "krea2"
+
+    uses_reference_latents = False
+
+    def _vision_view(self, image, block_state):
+        raise NotImplementedError
+
+    @property
+    def description(self) -> str:
+        return (
+            "Collect the reference images and size them for this mode's Qwen3-VL vision view, normalizing the "
+            "per-image style/subject strengths and masks."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="reference_images",
+                type_hint=PIL.Image.Image | list[PIL.Image.Image],
+                description="Reference image(s) to condition on. A single image or a list.",
+            ),
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description=(
+                    "How references condition the model: 'off' (Qwen3-VL vision path only, no LoRA), 'append' "
+                    "(vision + clean VAE tokens after the target, Ostris edit LoRA) or 'prepend' (vision + "
+                    "clean VAE source before the target, Identity-Edit LoRA)."
+                ),
+            ),
+            InputParam(
+                name="reference_style_strength",
+                type_hint=float | list[float],
+                default=1.0,
+                description=(
+                    "Per-image scale on the deepstack vision features (low-level texture/style). A scalar "
+                    "applies to every reference; 0 removes the channel."
+                ),
+            ),
+            InputParam(
+                name="reference_subject_strength",
+                type_hint=float | list[float],
+                default=1.0,
+                description=(
+                    "Per-image scale on the in-sequence image tokens (content/subject). A scalar applies to "
+                    "every reference; 0 removes the channel."
+                ),
+            ),
+            InputParam(
+                name="reference_masks",
+                type_hint=PIL.Image.Image | list[PIL.Image.Image],
+                description=(
+                    "Optional per-image masks restricting each reference to part of the picture: painted is "
+                    "the part to use. `reference_mask_mode` decides what happens to the rest. Index-aligned "
+                    "to `reference_images`; use `None` for slots without a mask."
+                ),
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(name="reference_mode", type_hint=str, description="The normalized reference mode."),
+            OutputParam(
+                name="vision_reference_images",
+                type_hint=list,
+                description="References sized for the Qwen3-VL vision path (empty when there are none).",
+            ),
+            OutputParam(
+                name="vae_reference_images",
+                type_hint=list,
+                description="References for the VAE reference-latent channel (empty outside the latent modes).",
+            ),
+            OutputParam(
+                name="reference_style_strengths", type_hint=list, description="One style scale per reference."
+            ),
+            OutputParam(
+                name="reference_subject_strengths", type_hint=list, description="One subject scale per reference."
+            ),
+            OutputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        images = [image for image in _as_list(block_state.reference_images) if image is not None]
+
+        block_state.reference_mode = self.reference_mode
+        block_state.reference_style_strengths = _broadcast_per_image(
+            block_state.reference_style_strength, len(images), 1.0
+        )
+        block_state.reference_subject_strengths = _broadcast_per_image(
+            block_state.reference_subject_strength, len(images), 1.0
+        )
+        masks = _as_list(block_state.reference_masks)
+        block_state.reference_masks = [masks[i] if i < len(masks) else None for i in range(len(images))]
+
+        block_state.vision_reference_images = [self._vision_view(image, block_state) for image in images]
+        block_state.vae_reference_images = [_to_rgb(image) for image in images] if self.uses_reference_latents else []
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Krea2VisionReferenceImagesStep(Krea2ReferenceImagesStep):
+    """Reference collector for the vision-only mode: a ~1 MP Qwen3-VL view, no VAE latent channel."""
+
+    reference_mode = REFERENCE_MODE_OFF
+    uses_reference_latents = False
+
+    @property
+    def description(self) -> str:
+        return "Collect the references and give each a ~1 MP Qwen3-VL view. Vision-only: no VAE latent channel."
+
+    def _vision_view(self, image, block_state):
+        return _cap_area(image, VISION_MAX_PIXELS)
+
+
+# auto_docstring
+class Krea2AppendReferenceImagesStep(Krea2ReferenceImagesStep):
+    """Reference collector for the Ostris `append` mode: a coarse ~384px Qwen3-VL view plus the VAE channel."""
+
+    reference_mode = REFERENCE_MODE_APPEND
+    uses_reference_latents = True
+
+    @property
+    def description(self) -> str:
+        return (
+            "Collect the references with a coarse ~384px Qwen3-VL view; the detail rides the VAE "
+            "reference-latent channel instead."
+        )
+
+    def _vision_view(self, image, block_state):
+        return _cap_area(image, VISION_EDIT_MAX_PIXELS)
+
+
+# auto_docstring
+class Krea2PrependReferenceImagesStep(Krea2ReferenceImagesStep):
+    """
+    Reference collector for the Identity-Edit `prepend` mode: a `grounding_px` longest-side Qwen3-VL view plus
+    the VAE channel.
+    """
+
+    reference_mode = REFERENCE_MODE_PREPEND
+    uses_reference_latents = True
+
+    @property
+    def description(self) -> str:
+        return (
+            "Collect the references with a `grounding_px` longest-side Qwen3-VL view, plus the images for the "
+            "VAE reference-latent channel."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(
+                name="grounding_px",
+                type_hint=int,
+                default=DEFAULT_GROUNDING_PX,
+                description=(
+                    "Longest-side cap (px) on the Qwen3-VL view; the identity-vs-adherence dial. Lower is a "
+                    "stronger edit, higher is stronger identity, 0 means native resolution."
+                ),
+            )
+        ]
+
+    def _vision_view(self, image, block_state):
+        return _cap_longest_side(image, int(block_state.grounding_px or 0))
+
+
+# auto_docstring
+class Krea2ReferenceTextEncoderStep(Krea2TextEncoderStep):
+    """
+    Text encoder step that feeds the reference images through Qwen3-VL's vision tower alongside the prompt,
+    keeping Krea 2's descriptor system template and 12-layer tap but inserting vision placeholders in the user
+    turn. Falls back to the stock text-only encode when there are no references.
+    """
+
+    model_name = "krea2"
+
+    # Identity-Edit trains its unconditional as the same source image with an empty instruction, so only that
+    # mode grounds the negative branch.
+    grounds_negative = False
+
+    def __init__(self):
+        super().__init__()
+        self._processor = None
+        self._processor_tokenizer = None
+
+    @property
+    def description(self) -> str:
+        return (
+            "Text encoder step that feeds the reference images through Qwen3-VL's vision tower alongside the "
+            "prompt, scaling each reference's vision features by its per-image style/subject strengths and "
+            "keeping only what a painted mask selects. Falls back to the stock text-only encode with no "
+            "references."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return super().expected_components + [
+            ComponentSpec("processor", Qwen3VLProcessor, description="The Qwen3-VL vision processor.")
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return super().inputs + [
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description="The normalized reference mode from the reference-images step.",
+            ),
+            InputParam(
+                name="vision_reference_images",
+                type_hint=list,
+                description="References sized for the Qwen3-VL vision path.",
+            ),
+            InputParam(name="reference_style_strengths", type_hint=list, description="One style scale per reference."),
+            InputParam(
+                name="reference_subject_strengths", type_hint=list, description="One subject scale per reference."
+            ),
+            InputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."),
+            InputParam(
+                name="reference_mask_mode",
+                type_hint=str,
+                default=REFERENCE_MASK_MODE_EXCLUDE_BLANK,
+                description=(
+                    "What a painted mask does to the vision path: 'exclude_blank' (default, mask the vision "
+                    "attention and blank the masked-out tokens), 'exclude' (mask the attention only, leaving "
+                    "the masked-out tokens in the sequence) or 'deemphasize' (blank the tokens only, which "
+                    "attenuates the region rather than removing it)."
+                ),
+            ),
+        ]
+
+    def _get_processor(self, components):
+        """The declared `processor` component, or one derived from the text encoder when the repo has none."""
+        if components.processor is not None:
+            return components.processor
+        tokenizer = components.tokenizer
+        if self._processor is None or self._processor_tokenizer is not tokenizer:
+            self._processor = _build_reference_processor(components.text_encoder, tokenizer)
+            self._processor_tokenizer = tokenizer
+        return self._processor
+
+    def _image_prompt(self, images) -> str:
+        """The vision-token layout for this mode's training template."""
+        raise NotImplementedError
+
+    def _encode_prompt_with_vision(
+        self, components, prompts, images, style_strengths, subject_strengths, masks, device, mask_mode
+    ):
+        """Encode `prompts` with `images` in the user turn, returning `(hidden_states, attention_mask)`.
+
+        Unlike the text-only path this is variable length: the processor expands each `<|image_pad|>` to the
+        image's token count and right-pads the batch. The vision-token layout comes from `_image_prompt`.
+        """
+        processor = self._get_processor(components)
+        text_encoder = components.text_encoder
+        prefix_idx = _PROMPT_TEMPLATE_ENCODE_START_IDX
+
+        image_prompt = self._image_prompt(images)
+
+        # The stock prefix/suffix keep the leading block byte-identical to the text-only path, so the same
+        # `prefix_idx` drop still holds; the vision tokens sit after it and survive.
+        text = [
+            _PROMPT_TEMPLATE_ENCODE_PREFIX + image_prompt + (p or "") + _PROMPT_TEMPLATE_ENCODE_SUFFIX for p in prompts
+        ]
+        # Every prompt in the batch gets its own copy of the reference images, in placeholder order.
+        model_inputs = processor(text=text, images=list(images) * len(prompts), padding=True, return_tensors="pt")
+        model_inputs = model_inputs.to(device)
+
+        # Skip the `get_image_features` wrap entirely when every strength is 1.0 and nothing is masked, so the
+        # plain full-strength path stays byte-identical to an unwrapped encode.
+        repeats = len(prompts)
+        style = list(style_strengths) * repeats
+        subject = list(subject_strengths) * repeats
+        masks = list(masks) * repeats
+
+        has_pixels = model_inputs.get("pixel_values") is not None
+        needs_scaling = any(float(s) != 1.0 for s in style) or any(float(s) != 1.0 for s in subject)
+        has_masks = any(m is not None for m in masks)
+
+        # The mode selects which of the two independent layers a mask acts on.
+        mask_attention = has_masks and mask_mode != REFERENCE_MASK_MODE_DEEMPHASIZE
+        zero_masked_tokens = mask_mode != REFERENCE_MASK_MODE_EXCLUDE
+
+        keep_vectors = []
+        if has_pixels and zero_masked_tokens and (needs_scaling or has_masks):
+            keep_vectors = _build_keep_vectors(
+                masks, model_inputs.get("image_grid_thw"), processor.image_processor.merge_size
+            )
+        apply_scaling = has_pixels and (needs_scaling or any(k is not None for k in keep_vectors))
+
+        # Only the attention layer can keep unpainted content out of the reference; the output tokens the
+        # wrapper zeroes are whole-image summaries, so zeroing them attenuates but never excludes.
+        patch_keeps = []
+        if has_pixels and mask_attention:
+            patch_keeps = build_patch_keeps(
+                masks,
+                model_inputs.get("image_grid_thw"),
+                processor.image_processor.merge_size,
+                _mask_to_keep_grid,
+            )
+
+        original_get_image_features = None
+        if apply_scaling:
+            original_get_image_features = text_encoder.get_image_features
+            text_encoder.get_image_features = _scaled_get_image_features(
+                original_get_image_features, style, subject, processor.image_processor.merge_size, keep_vectors
+            )
+        try:
+            # No `position_ids`: with `pixel_values` present Qwen3-VL derives its mRoPE positions internally.
+            with vision_key_masks(patch_keeps):
+                outputs = text_encoder(
+                    input_ids=model_inputs["input_ids"],
+                    attention_mask=model_inputs["attention_mask"],
+                    pixel_values=model_inputs.get("pixel_values"),
+                    image_grid_thw=model_inputs.get("image_grid_thw"),
+                    mm_token_type_ids=model_inputs.get("mm_token_type_ids"),
+                    output_hidden_states=True,
+                )
+        finally:
+            if original_get_image_features is not None:
+                text_encoder.get_image_features = original_get_image_features
+
+        hidden_states = torch.stack([outputs.hidden_states[i] for i in KREA2_TEXT_ENCODER_SELECT_LAYERS], dim=2)
+        attention_mask = model_inputs["attention_mask"].bool()
+        return hidden_states[:, prefix_idx:], attention_mask[:, prefix_idx:]
+
+    @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)
+        images = list(block_state.vision_reference_images or [])
+        mask_mode = _coerce_reference_mask_mode(block_state.reference_mask_mode)
+
+        if images:
+            block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt_with_vision(
+                components,
+                prompts,
+                images,
+                block_state.reference_style_strengths or [],
+                block_state.reference_subject_strengths or [],
+                block_state.reference_masks or [],
+                device,
+                mask_mode,
+            )
+        else:
+            block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt(
+                components, prompts, block_state.max_sequence_length, device
+            )
+
+        block_state.negative_prompt_embeds = None
+        block_state.negative_prompt_embeds_mask = None
+        if components.requires_unconditional_embeds:
+            negative_prompt = block_state.negative_prompt
+            if negative_prompt is None:
+                negative_prompt = ""
+            if isinstance(negative_prompt, str):
+                negative_prompt = [negative_prompt] * len(prompts)
+
+            if self.grounds_negative and images:
+                # The negative MUST get the same masks as the positive: an unmasked negative sees the whole
+                # reference, so (cond - uncond) pushes the unpainted region away instead of ignoring it.
+                # Strengths stay at full -- that is the training condition.
+                block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = (
+                    self._encode_prompt_with_vision(
+                        components,
+                        negative_prompt,
+                        images,
+                        [],
+                        [],
+                        block_state.reference_masks or [],
+                        device,
+                        mask_mode,
+                    )
+                )
+            else:
+                block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._encode_prompt(
+                    components, negative_prompt, block_state.max_sequence_length, device
+                )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Krea2VisionReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep):
+    """
+    Text encoder for the vision-only mode. A single reference goes in as a bare vision block; several are
+    labelled "Picture N:". The negative branch stays text-only.
+    """
+
+    grounds_negative = False
+
+    @property
+    def description(self) -> str:
+        return "Text encoder for the vision-only mode: references enter the prompt through Qwen3-VL's vision tower."
+
+    def _image_prompt(self, images) -> str:
+        if len(images) > 1:
+            return "".join(f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>" for i in range(len(images)))
+        return "<|vision_start|><|image_pad|><|vision_end|>"
+
+
+# auto_docstring
+class Krea2AppendReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep):
+    """
+    Text encoder for the Ostris `append` mode, which labels every reference "Picture N:" even when there is only
+    one. The negative branch stays text-only.
+    """
+
+    grounds_negative = False
+
+    @property
+    def description(self) -> str:
+        return 'Text encoder for the Ostris append mode: every reference is labelled "Picture N:".'
+
+    def _image_prompt(self, images) -> str:
+        return "".join(f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>" for i in range(len(images)))
+
+
+# auto_docstring
+class Krea2PrependReferenceTextEncoderStep(Krea2ReferenceTextEncoderStep):
+    """Text encoder for the Identity-Edit `prepend` mode: unlabelled vision blocks, grounded negative branch."""
+
+    grounds_negative = True
+
+    @property
+    def description(self) -> str:
+        return (
+            "Text encoder for the Identity-Edit prepend mode: bare vision blocks and a negative branch grounded "
+            "on the same source, as the LoRA was trained."
+        )
+
+    def _image_prompt(self, images) -> str:
+        return "".join("<|vision_start|><|image_pad|><|vision_end|>" for _ in images)
+
+
+# auto_docstring
+class Krea2ReferenceTextInputsStep(ModularPipelineBlocks):
+    """
+    Input step that determines `batch_size`/`dtype` from the per-prompt `prompt_embeds` and replicates the text
+    conditioning (and the negative branch) to `batch_size * num_images_per_prompt`. Unlike the stock step it
+    lets the two branches carry different sequence lengths.
+    """
+
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Input step that determines `batch_size`/`dtype` and batch-expands the text conditioning, allowing "
+            "the positive and negative branches to carry different sequence lengths (the positive gains vision "
+            "tokens the text-only negative does not have)."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_images_per_prompt", default=1),
+            InputParam.template("prompt_embeds"),
+            InputParam.template("prompt_embeds_mask"),
+            InputParam.template("negative_prompt_embeds"),
+            InputParam.template("negative_prompt_embeds_mask"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="batch_size",
+                type_hint=int,
+                description="Effective batch size (num prompts * num_images_per_prompt).",
+            ),
+            OutputParam(name="dtype", type_hint=torch.dtype, description="The dtype of the text features."),
+            OutputParam(name="prompt_embeds", type_hint=torch.Tensor, description="Text features, batch-expanded."),
+            OutputParam(name="prompt_embeds_mask", type_hint=torch.Tensor, description="Text mask, batch-expanded."),
+            OutputParam(
+                name="negative_prompt_embeds",
+                type_hint=torch.Tensor,
+                description="Negative text features, batch-expanded.",
+            ),
+            OutputParam(
+                name="negative_prompt_embeds_mask",
+                type_hint=torch.Tensor,
+                description="Negative text mask, batch-expanded.",
+            ),
+        ]
+
+    @staticmethod
+    def _expand(embeds, mask, num_images_per_prompt):
+        prompt_batch, seq_len, num_layers, dim = embeds.shape
+        n = num_images_per_prompt
+        embeds = embeds.repeat(1, n, 1, 1).view(prompt_batch * n, seq_len, num_layers, dim)
+        mask = mask.repeat(1, n).view(prompt_batch * n, seq_len)
+        return embeds, mask
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        n = block_state.num_images_per_prompt
+        block_state.dtype = block_state.prompt_embeds.dtype
+        block_state.batch_size = block_state.prompt_embeds.shape[0] * n
+
+        block_state.prompt_embeds, block_state.prompt_embeds_mask = self._expand(
+            block_state.prompt_embeds, block_state.prompt_embeds_mask, n
+        )
+        if block_state.negative_prompt_embeds is not None:
+            block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = self._expand(
+                block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask, n
+            )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Krea2ReferenceLatentsStep(ModularPipelineBlocks):
+    """
+    VAE-encode the reference images into clean latent tokens for the `append`/`prepend` conditioning paths.
+    Each reference is encoded at up to ~1 MP with its own aspect ratio, normalized with the VAE's per-channel
+    statistics, patch-packed exactly like the noise latents and given Kontext-style rotary coordinates: the
+    i-th reference sits on rotary frame axis `i + 1` with its own y/x grid. Reference masks reach this channel
+    only when `mask_reference_latents` is set. A no-op in `off` mode or with no references.
+    """
+
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "VAE-encode the reference images into clean, patch-packed latent tokens with frame-axis rotary "
+            "coordinates for the `append`/`prepend` conditioning paths. A no-op in `off` mode."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("vae", AutoencoderKLQwenImage)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description="The normalized reference mode from the reference-images step.",
+            ),
+            InputParam(
+                name="vae_reference_images",
+                type_hint=list,
+                description="References for the VAE reference-latent channel.",
+            ),
+            InputParam(name="reference_masks", type_hint=list, description="One mask (or None) per reference."),
+            InputParam(
+                name="mask_reference_latents",
+                type_hint=bool,
+                default=False,
+                description=(
+                    "Apply the reference masks to this channel too, on top of the vision tokens they always "
+                    "mask. Off (default) is vision-only masking, so the edit LoRA still sees the whole "
+                    "reference at full detail -- which is why a masked edit can show hints of the unpainted "
+                    "parts. On drops the unpainted latent patch tokens from the sequence entirely."
+                ),
+            ),
+            InputParam.template("dtype"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="reference_latents",
+                type_hint=torch.Tensor,
+                description="Clean packed reference tokens (1, reference_seq_len, in_channels), or None.",
+            ),
+            OutputParam(
+                name="reference_position_ids",
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the reference tokens (reference_seq_len, 3), or None.",
+            ),
+            OutputParam(
+                name="reference_seq_len", type_hint=int, description="Number of reference tokens (0 when unused)."
+            ),
+        ]
+
+    @staticmethod
+    def _to_snapped_tensor(image, snap: int) -> torch.Tensor:
+        """Reference image -> normalized `(1, 3, 1, H, W)` tensor for the temporal VAE.
+
+        Aspect-preserving downscale to fit `REFERENCE_LATENTS_MAX_PIXELS` (never upscaled), then snap each side
+        to a multiple of `snap` so the latent grid is patchifiable.
+        """
+        image = _to_rgb(image)
+        width, height = image.width, image.height
+        if height * width > REFERENCE_LATENTS_MAX_PIXELS:
+            ratio = height / width
+            new_height = math.sqrt(REFERENCE_LATENTS_MAX_PIXELS * ratio)
+            new_width = math.sqrt(REFERENCE_LATENTS_MAX_PIXELS / ratio)
+        else:
+            new_height, new_width = float(height), float(width)
+        new_height = max(snap, int(round(new_height / snap)) * snap)
+        new_width = max(snap, int(round(new_width / snap)) * snap)
+
+        array = np.asarray(image, dtype=np.float32) / 255.0
+        tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) * 2.0 - 1.0
+        if (new_height, new_width) != (height, width):
+            tensor = F.interpolate(tensor, size=(new_height, new_width), mode="bilinear", align_corners=False)
+        return tensor.unsqueeze(2)
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        images = list(block_state.vae_reference_images or [])
+        if not images:
+            raise ValueError(
+                "The append and prepend modes need at least one reference image; pass `reference_images` or "
+                "switch to the vision-only mode."
+            )
+
+        block_state.reference_latents = None
+        block_state.reference_position_ids = None
+        block_state.reference_seq_len = 0
+
+        vae = components.vae
+        patch_size = components.patch_size
+        snap = components.vae_scale_factor * patch_size
+
+        z_dim = vae.config.z_dim
+        latents_mean = torch.tensor(vae.config.latents_mean).view(1, z_dim, 1, 1, 1)
+        latents_std = torch.tensor(vae.config.latents_std).view(1, z_dim, 1, 1, 1)
+
+        masks = list(block_state.reference_masks or [])
+        mask_latents = bool(block_state.mask_reference_latents)
+
+        tokens = []
+        position_ids = []
+        for i, image in enumerate(images):
+            pixels = self._to_snapped_tensor(image, snap).to(vae.device, vae.dtype)
+            raw = vae.encode(pixels).latent_dist.mode()
+            mean = latents_mean.to(raw.device, raw.dtype)
+            std = latents_std.to(raw.device, raw.dtype)
+            latents = ((raw - mean) / std)[:, :, 0]  # drop the temporal axis -> (1, C, lat_h, lat_w)
+
+            image_tokens = _pack_latents(latents, patch_size).to(block_state.dtype)
+            _, _, latent_height, latent_width = latents.shape
+            grid_height, grid_width = latent_height // patch_size, latent_width // patch_size
+            image_position_ids = _reference_position_ids(i, grid_height, grid_width)
+
+            if mask_latents:
+                keep = _mask_to_keep_vector(masks[i] if i < len(masks) else None, grid_height, grid_width)
+                if keep is not None:
+                    # Dropped rather than zeroed: a zero token is the mean latent (flat mid-grey), which the
+                    # model can still attend to and reproduce. Position ids go with them, so the kept tokens
+                    # keep their true coordinates.
+                    index = torch.nonzero(keep > 0.5, as_tuple=False).squeeze(-1)
+                    if index.numel() == 0:
+                        continue
+                    image_tokens = image_tokens[:, index]
+                    image_position_ids = image_position_ids[index]
+
+            tokens.append(image_tokens)
+            position_ids.append(image_position_ids)
+
+        block_state.reference_latents = torch.cat(tokens, dim=1)
+        block_state.reference_position_ids = torch.cat(position_ids, dim=0)
+        block_state.reference_seq_len = int(block_state.reference_latents.shape[1])
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Krea2ReferencePreparePositionIdsStep(ModularPipelineBlocks):
+    """
+    Build the rotary position ids for both guidance branches, which carry different text lengths once the
+    positive gains vision tokens. Reference-latent coordinates are spliced in before the target for `prepend`
+    and after it for `append`.
+    """
+
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Build the rotary position ids for both guidance branches (the branches carry different text "
+            "lengths once the positive gains vision tokens), splicing in the reference-latent coordinates "
+            "before the target for `prepend` and after it for `append`."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("height", default=1024),
+            InputParam.template("width", default=1024),
+            InputParam.template("prompt_embeds"),
+            InputParam.template("negative_prompt_embeds"),
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description="The normalized reference mode from the reference-images step.",
+            ),
+            InputParam(
+                name="reference_position_ids",
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the reference tokens, when the latent modes are active.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="position_ids",
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the conditional branch.",
+            ),
+            OutputParam(
+                name="negative_position_ids",
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the unconditional branch.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        device = components._execution_device
+        patch_size = components.patch_size
+        grid_h = block_state.height // (components.vae_scale_factor * patch_size)
+        grid_w = block_state.width // (components.vae_scale_factor * patch_size)
+        reference_ids = block_state.reference_position_ids
+        if reference_ids is not None:
+            reference_ids = reference_ids.to(device)
+
+        def build(text_seq_len: int) -> torch.Tensor:
+            ids = Krea2PreparePositionIdsStep.prepare_position_ids(text_seq_len, grid_h, grid_w, device)
+            if reference_ids is None or self.reference_mode == REFERENCE_MODE_OFF:
+                return ids
+            if self.reference_mode == REFERENCE_MODE_PREPEND:
+                # [text | reference | target]
+                return torch.cat([ids[:text_seq_len], reference_ids, ids[text_seq_len:]], dim=0)
+            # [text | target | reference]
+            return torch.cat([ids, reference_ids], dim=0)
+
+        text_seq_len = block_state.prompt_embeds.shape[1]
+        block_state.position_ids = build(text_seq_len)
+
+        negative_prompt_embeds = block_state.negative_prompt_embeds
+        if negative_prompt_embeds is None or negative_prompt_embeds.shape[1] == text_seq_len:
+            block_state.negative_position_ids = block_state.position_ids
+        else:
+            block_state.negative_position_ids = build(negative_prompt_embeds.shape[1])
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class Krea2ReferenceLoopDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Within the denoising loop: concatenate the clean reference tokens onto the noisy latents in the "
+            "order the mode dictates, run the `transformer` per guidance branch with that branch's own "
+            "position ids, and slice the reference rows back off the prediction. Compose into "
+            "`Krea2ReferenceDenoiseStep`."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "guider",
+                ClassifierFreeGuidance,
+                # Krea 2 uses cond-anchored CFG (`cond + scale * (cond - uncond)`), which is the
+                # `use_original_formulation` branch.
+                config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}),
+                default_creation_method="from_config",
+            ),
+            ComponentSpec("transformer", Krea2Transformer2DModel),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+            InputParam.template("num_inference_steps", required=True),
+            InputParam.template("prompt_embeds"),
+            InputParam.template("prompt_embeds_mask"),
+            InputParam.template("negative_prompt_embeds"),
+            InputParam.template("negative_prompt_embeds_mask"),
+            InputParam(
+                name="position_ids",
+                required=True,
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the conditional branch.",
+            ),
+            InputParam(
+                name="negative_position_ids",
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the unconditional branch.",
+            ),
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description="The normalized reference mode from the reference-images step.",
+            ),
+            InputParam(
+                name="reference_latents",
+                type_hint=torch.Tensor,
+                description="Clean packed reference tokens, when the latent modes are active.",
+            ),
+            InputParam.template("attention_kwargs"),
+        ]
+
+    @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)
+        target_seq_len = latents.shape[1]
+
+        latents, reference_seq_len = _concat_reference_latents(
+            latents, block_state.reference_latents, self.reference_mode, transformer.dtype
+        )
+
+        negative_position_ids = block_state.negative_position_ids
+        if negative_position_ids is None:
+            negative_position_ids = block_state.position_ids
+
+        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,
+            ),
+            # A tuple is indexed per guidance branch, so each pass gets coordinates for its own text length.
+            "position_ids": (block_state.position_ids, negative_position_ids),
+        }
+
+        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}
+
+            handles = []
+            if reference_seq_len and self.reference_mode == REFERENCE_MODE_APPEND:
+                # The split index is branch-dependent: the conditional text block carries the vision tokens.
+                text_seq_len = cond_kwargs["encoder_hidden_states"].shape[1]
+                handles = _register_zero_time_reference_hooks(
+                    transformer,
+                    timestep,
+                    batch_size=latents.shape[0],
+                    split=text_seq_len + target_seq_len,
+                    total=text_seq_len + target_seq_len + reference_seq_len,
+                )
+            try:
+                noise_pred = transformer(
+                    hidden_states=latents,
+                    timestep=timestep,
+                    attention_kwargs=block_state.attention_kwargs,
+                    return_dict=False,
+                    **cond_kwargs,
+                )[0]
+            finally:
+                for handle in handles:
+                    handle.remove()
+
+            guider_state_batch.noise_pred = _slice_reference_rows(
+                noise_pred, self.reference_mode, reference_seq_len, target_seq_len
+            )
+            components.guider.cleanup_models(transformer)
+
+        block_state.noise_pred = components.guider(guider_state).pred
+        return components, block_state
+
+
+# auto_docstring
+class Krea2ReferenceDenoiseStep(Krea2DenoiseLoopWrapper):
+    """
+    Denoising loop that iteratively denoises the packed image latents over `timesteps` with reference
+    conditioning: the clean reference tokens are concatenated in the order the mode dictates, each guidance
+    branch runs with its own position ids, and the reference rows are sliced back off the prediction before the
+    scheduler step.
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2LoopBeforeDenoiser, Krea2ReferenceLoopDenoiser, Krea2LoopAfterDenoiser]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoising loop with reference conditioning: concatenates the clean reference tokens, runs each "
+            "guidance branch with its own position ids, and slices the reference rows back off the prediction. "
+            "Identical to the stock loop when no reference latents are active."
+        )
+
+
+# =====================================================================================================
+# Per-mode leaf blocks
+# =====================================================================================================
+
+
+# auto_docstring
+class Krea2VisionReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep):
+    """Rotary coordinates for the vision-only mode: per-branch, with no reference rows to splice in."""
+
+    reference_mode = REFERENCE_MODE_OFF
+
+    @property
+    def description(self) -> str:
+        return "Per-branch rotary coordinates for the [text | image] sequence (vision-only: no reference rows)."
+
+
+# auto_docstring
+class Krea2AppendReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep):
+    """Rotary coordinates for the Ostris `append` mode: the reference rows go after the target."""
+
+    reference_mode = REFERENCE_MODE_APPEND
+
+    @property
+    def description(self) -> str:
+        return "Per-branch rotary coordinates with the reference rows appended after the target."
+
+
+# auto_docstring
+class Krea2PrependReferencePositionIdsStep(Krea2ReferencePreparePositionIdsStep):
+    """Rotary coordinates for the Identity-Edit `prepend` mode: the reference rows go before the target."""
+
+    reference_mode = REFERENCE_MODE_PREPEND
+
+    @property
+    def description(self) -> str:
+        return "Per-branch rotary coordinates with the reference rows spliced in before the target."
+
+
+class Krea2VisionReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_OFF
+
+    @property
+    def description(self) -> str:
+        return (
+            "Within the denoising loop: run the `transformer` per guidance branch with that branch's own "
+            "position ids. Vision-only, so there are no reference tokens in the sequence."
+        )
+
+
+class Krea2AppendReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_APPEND
+
+    @property
+    def description(self) -> str:
+        return (
+            "Within the denoising loop: append the clean reference tokens after the noisy target, modulate them "
+            "at flow time t=0, and slice them back off the prediction."
+        )
+
+
+class Krea2PrependReferenceLoopDenoiser(Krea2ReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_PREPEND
+
+    @property
+    def description(self) -> str:
+        return (
+            "Within the denoising loop: prepend the clean reference tokens before the noisy target under the "
+            "plain uniform modulation, and slice them back off the prediction."
+        )
+
+
+# auto_docstring
+class Krea2VisionReferenceDenoiseStep(Krea2DenoiseLoopWrapper):
+    """Denoising loop for the vision-only mode, with symmetric CFG."""
+
+    model_name = "krea2"
+    block_classes = [Krea2LoopBeforeDenoiser, Krea2VisionReferenceLoopDenoiser, Krea2LoopAfterDenoiser]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return "Denoising loop for the vision-only mode, with per-branch position ids."
+
+
+# auto_docstring
+class Krea2AppendReferenceDenoiseStep(Krea2DenoiseLoopWrapper):
+    """Denoising loop for the Ostris `append` mode: reference tokens at the tail, modulated at t=0."""
+
+    model_name = "krea2"
+    block_classes = [Krea2LoopBeforeDenoiser, Krea2AppendReferenceLoopDenoiser, Krea2LoopAfterDenoiser]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return "Denoising loop for the Ostris append mode: clean reference tokens at the tail, modulated at t=0."
+
+
+# auto_docstring
+class Krea2PrependReferenceDenoiseStep(Krea2DenoiseLoopWrapper):
+    """Denoising loop for the Identity-Edit `prepend` mode: the clean source sits before the target."""
+
+    model_name = "krea2"
+    block_classes = [Krea2LoopBeforeDenoiser, Krea2PrependReferenceLoopDenoiser, Krea2LoopAfterDenoiser]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return "Denoising loop for the Identity-Edit prepend mode: the clean source before the noisy target."
+
+
+class _Krea2TurboTextEncoderMixin:
+    """Drops the negative branch from a reference text encoder for the distilled checkpoint."""
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [spec for spec in super().expected_components if spec.name != "guider"]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [param for param in super().inputs if param.name != "negative_prompt"]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            param
+            for param in super().intermediate_outputs
+            if param.name not in ("negative_prompt_embeds", "negative_prompt_embeds_mask")
+        ]
+
+    @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)
+        images = list(block_state.vision_reference_images or [])
+        mask_mode = _coerce_reference_mask_mode(block_state.reference_mask_mode)
+
+        if images:
+            block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt_with_vision(
+                components,
+                prompts,
+                images,
+                block_state.reference_style_strengths or [],
+                block_state.reference_subject_strengths or [],
+                block_state.reference_masks or [],
+                device,
+                mask_mode,
+            )
+        else:
+            block_state.prompt_embeds, block_state.prompt_embeds_mask = self._encode_prompt(
+                components, prompts, block_state.max_sequence_length, device
+            )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# auto_docstring
+class Krea2TurboVisionReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2VisionReferenceTextEncoderStep):
+    """Vision-only text encoder for the distilled checkpoint: no negative branch, no guider."""
+
+
+# auto_docstring
+class Krea2TurboAppendReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2AppendReferenceTextEncoderStep):
+    """Ostris append text encoder for the distilled checkpoint: no negative branch, no guider."""
+
+
+# auto_docstring
+class Krea2TurboPrependReferenceTextEncoderStep(_Krea2TurboTextEncoderMixin, Krea2PrependReferenceTextEncoderStep):
+    """Identity-Edit prepend text encoder for the distilled checkpoint: no negative branch, no guider."""
+
+
+class Krea2TurboReferenceLoopDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Within the denoising loop: run the `transformer` on the conditional text features with the clean "
+            "reference tokens spliced in, then slice them back off. The distilled checkpoint runs without "
+            "classifier-free guidance, so there is no negative branch or guider."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [ComponentSpec("transformer", Krea2Transformer2DModel)]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(name="latents", required=True, type_hint=torch.Tensor, description="Packed image latents."),
+            InputParam.template("prompt_embeds"),
+            InputParam.template("prompt_embeds_mask"),
+            InputParam(
+                name="position_ids",
+                required=True,
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the [text | image] sequence.",
+            ),
+            InputParam(
+                name="reference_mode",
+                type_hint=str,
+                default=REFERENCE_MODE_OFF,
+                description="The normalized reference mode from the reference-images step.",
+            ),
+            InputParam(
+                name="reference_latents",
+                type_hint=torch.Tensor,
+                description="Clean packed reference tokens, when the latent modes are active.",
+            ),
+            InputParam.template("attention_kwargs"),
+        ]
+
+    @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)
+        target_seq_len = latents.shape[1]
+
+        latents, reference_seq_len = _concat_reference_latents(
+            latents, block_state.reference_latents, self.reference_mode, transformer.dtype
+        )
+
+        handles = []
+        if reference_seq_len and self.reference_mode == REFERENCE_MODE_APPEND:
+            text_seq_len = block_state.prompt_embeds.shape[1]
+            handles = _register_zero_time_reference_hooks(
+                transformer,
+                timestep,
+                batch_size=latents.shape[0],
+                split=text_seq_len + target_seq_len,
+                total=text_seq_len + target_seq_len + reference_seq_len,
+            )
+        try:
+            noise_pred = transformer(
+                hidden_states=latents,
+                timestep=timestep,
+                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]
+        finally:
+            for handle in handles:
+                handle.remove()
+
+        block_state.noise_pred = _slice_reference_rows(
+            noise_pred, self.reference_mode, reference_seq_len, target_seq_len
+        )
+        return components, block_state
+
+
+# auto_docstring
+class Krea2TurboReferenceDenoiseStep(Krea2DenoiseLoopWrapper):
+    """
+    Denoising loop for the distilled Krea 2 Turbo checkpoint with reference conditioning: the clean reference
+    tokens are spliced into the sequence and sliced back off the prediction. No classifier-free guidance.
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2LoopBeforeDenoiser, Krea2TurboReferenceLoopDenoiser, Krea2LoopAfterDenoiser]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoising loop for the distilled Krea 2 Turbo checkpoint with reference conditioning, without "
+            "classifier-free guidance."
+        )
+
+
+class Krea2TurboVisionReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_OFF
+
+
+class Krea2TurboAppendReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_APPEND
+
+
+class Krea2TurboPrependReferenceLoopDenoiser(Krea2TurboReferenceLoopDenoiser):
+    reference_mode = REFERENCE_MODE_PREPEND
+
+
+def _turbo_denoise_step(name, denoiser, summary):
+    """Build a guidance-free denoise loop for one mode."""
+    return type(
+        name,
+        (Krea2DenoiseLoopWrapper,),
+        {
+            "__doc__": summary,
+            "model_name": "krea2",
+            "block_classes": [Krea2LoopBeforeDenoiser, denoiser, Krea2LoopAfterDenoiser],
+            "block_names": ["before_denoiser", "denoiser", "after_denoiser"],
+            "description": property(lambda self, _s=summary: _s),
+        },
+    )
+
+
+Krea2TurboVisionReferenceDenoiseStep = _turbo_denoise_step(
+    "Krea2TurboVisionReferenceDenoiseStep",
+    Krea2TurboVisionReferenceLoopDenoiser,
+    "Guidance-free denoising loop for the vision-only mode on the distilled checkpoint.",
+)
+Krea2TurboAppendReferenceDenoiseStep = _turbo_denoise_step(
+    "Krea2TurboAppendReferenceDenoiseStep",
+    Krea2TurboAppendReferenceLoopDenoiser,
+    "Guidance-free denoising loop for the Ostris append mode on the distilled checkpoint.",
+)
+Krea2TurboPrependReferenceDenoiseStep = _turbo_denoise_step(
+    "Krea2TurboPrependReferenceDenoiseStep",
+    Krea2TurboPrependReferenceLoopDenoiser,
+    "Guidance-free denoising loop for the Identity-Edit prepend mode on the distilled checkpoint.",
+)
+
+
+# =====================================================================================================
+# Workflows: one flat blockset per mode, per checkpoint.
+# =====================================================================================================
+
+
+def _workflow(name, images, encoder, timesteps, position_ids, denoise, latents, summary):
+    core = InsertableDict(
+        [("input", Krea2ReferenceTextInputsStep())]
+        + ([("reference_latents", Krea2ReferenceLatentsStep())] if latents else [])
+        + [
+            ("prepare_latents", Krea2PrepareLatentsStep()),
+            ("set_timesteps", timesteps()),
+            ("prepare_position_ids", position_ids()),
+            ("denoise", denoise()),
+        ]
+    )
+    core_cls = type(
+        f"{name}CoreDenoiseStep",
+        (SequentialPipelineBlocks,),
+        {
+            "__doc__": f"Core denoising workflow: {summary}",
+            "model_name": "krea2",
+            "block_classes": list(core.values()),
+            "block_names": list(core.keys()),
+            "description": property(lambda self, _s=summary: f"Core denoising workflow: {_s}"),
+            "outputs": property(
+                lambda self: [
+                    OutputParam.template(
+                        "latents", description="The denoised packed latents (B, image_seq_len, in_channels)."
+                    )
+                ]
+            ),
+        },
+    )
+    return type(
+        name,
+        (SequentialPipelineBlocks,),
+        {
+            "__doc__": summary,
+            "model_name": "krea2",
+            "block_classes": [images, encoder, core_cls, Krea2DecodeStep],
+            "block_names": ["reference_images", "text_encoder", "denoise", "decode"],
+            "description": property(lambda self, _s=summary: _s),
+            "outputs": property(lambda self: [OutputParam.template("images")]),
+        },
+    )
+
+
+Krea2VisionReferenceWorkflow = _workflow(
+    "Krea2VisionReferenceWorkflow",
+    Krea2VisionReferenceImagesStep,
+    Krea2VisionReferenceTextEncoderStep,
+    Krea2SetTimestepsStep,
+    Krea2VisionReferencePositionIdsStep,
+    Krea2VisionReferenceDenoiseStep,
+    latents=False,
+    summary="vision-only references on the stock checkpoint, no LoRA needed.",
+)
+Krea2AppendReferenceWorkflow = _workflow(
+    "Krea2AppendReferenceWorkflow",
+    Krea2AppendReferenceImagesStep,
+    Krea2AppendReferenceTextEncoderStep,
+    Krea2SetTimestepsStep,
+    Krea2AppendReferencePositionIdsStep,
+    Krea2AppendReferenceDenoiseStep,
+    latents=True,
+    summary="vision plus clean VAE reference tokens after the target, for the Ostris edit LoRA.",
+)
+Krea2PrependReferenceWorkflow = _workflow(
+    "Krea2PrependReferenceWorkflow",
+    Krea2PrependReferenceImagesStep,
+    Krea2PrependReferenceTextEncoderStep,
+    Krea2SetTimestepsStep,
+    Krea2PrependReferencePositionIdsStep,
+    Krea2PrependReferenceDenoiseStep,
+    latents=True,
+    summary="vision plus the clean VAE source before the target, for the Identity-Edit LoRA.",
+)
+
+Krea2TurboVisionReferenceWorkflow = _workflow(
+    "Krea2TurboVisionReferenceWorkflow",
+    Krea2VisionReferenceImagesStep,
+    Krea2TurboVisionReferenceTextEncoderStep,
+    Krea2TurboSetTimestepsStep,
+    Krea2VisionReferencePositionIdsStep,
+    Krea2TurboVisionReferenceDenoiseStep,
+    latents=False,
+    summary="vision-only references on the distilled checkpoint.",
+)
+Krea2TurboAppendReferenceWorkflow = _workflow(
+    "Krea2TurboAppendReferenceWorkflow",
+    Krea2AppendReferenceImagesStep,
+    Krea2TurboAppendReferenceTextEncoderStep,
+    Krea2TurboSetTimestepsStep,
+    Krea2AppendReferencePositionIdsStep,
+    Krea2TurboAppendReferenceDenoiseStep,
+    latents=True,
+    summary="vision plus clean VAE reference tokens after the target, for the Ostris Turbo edit LoRA.",
+)
+Krea2TurboPrependReferenceWorkflow = _workflow(
+    "Krea2TurboPrependReferenceWorkflow",
+    Krea2PrependReferenceImagesStep,
+    Krea2TurboPrependReferenceTextEncoderStep,
+    Krea2TurboSetTimestepsStep,
+    Krea2PrependReferencePositionIdsStep,
+    Krea2TurboPrependReferenceDenoiseStep,
+    latents=True,
+    summary="vision plus the clean VAE source before the target, on the distilled checkpoint.",
+)
+
+
+# auto_docstring
+class Krea2ReferenceAutoBlocks(ConditionalPipelineBlocks):
+    """
+    Modular pipeline for Krea 2 with reference-image conditioning. `reference_mode` picks the workflow:
+    `"off"` (vision path only, stock checkpoint), `"append"` (Ostris edit LoRA) or `"prepend"` (Identity-Edit
+    LoRA). With no `reference_images` this is plain text-to-image.
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2VisionReferenceWorkflow,
+        Krea2AppendReferenceWorkflow,
+        Krea2PrependReferenceWorkflow,
+    ]
+    block_names = [REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND]
+    block_trigger_inputs = ["reference_mode"]
+    default_block_name = REFERENCE_MODE_OFF
+
+    @property
+    def description(self) -> str:
+        return (
+            "Krea 2 with reference-image conditioning: `reference_mode` selects the vision-only path, the "
+            "Ostris append path or the Identity-Edit prepend path."
+        )
+
+    def select_block(self, **kwargs) -> str:
+        return _coerce_reference_mode(kwargs.get("reference_mode"))
+
+
+# auto_docstring
+class Krea2TurboReferenceAutoBlocks(Krea2ReferenceAutoBlocks):
+    """
+    Modular pipeline for the distilled Krea 2 Turbo checkpoint with reference-image conditioning. The same three
+    modes on the distilled schedule and without classifier-free guidance, so it takes no negative prompt and
+    carries no guider.
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2TurboVisionReferenceWorkflow,
+        Krea2TurboAppendReferenceWorkflow,
+        Krea2TurboPrependReferenceWorkflow,
+    ]
+    block_names = [REFERENCE_MODE_OFF, REFERENCE_MODE_APPEND, REFERENCE_MODE_PREPEND]
+    block_trigger_inputs = ["reference_mode"]
+    default_block_name = REFERENCE_MODE_OFF
+
+    @property
+    def description(self) -> str:
+        return (
+            "Krea 2 Turbo with reference-image conditioning: the distilled schedule, no CFG, and the same "
+            "three reference modes."
+        )

krea2_vision_attention.py

@@ -0,0 +1,215 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Restrict Qwen3-VL's vision self-attention to the painted region of a Krea 2 reference image.
+
+Zeroing the vision tower's OUTPUT tokens cannot remove content from a reference. `Qwen3VLVisionModel`
+runs its full depth of self-attention over every patch of the image (`cu_seqlens` only separates one
+image from the next), so every output token is a contextual summary of the whole picture: dropping the
+tokens that cover a subject deletes its dedicated slots and nothing else. Removal requires masking
+inside the attention, so kept patches never attend to unpainted patches as keys.
+
+`Qwen3VLVisionAttention.forward` is patched to accept a per-image key mask, active only inside
+`vision_key_masks(...)`. With no active mask it defers to the original forward, so unmasked encodes are
+unchanged.
+
+Two constraints that are silent when broken:
+
+* The mask lives at **pre-merge patch resolution** (`grid_h x grid_w`) and the flattened patch sequence
+  is **not row-major**. The image processor groups patches by `merge_size x merge_size` window
+  (`permute(0, 2, 5, 3, 6, 1, 4, 7)` over `(batch, channel, gh/m, m, ps, gw/m, m, ps)`), giving
+  sequence order `(h_block, w_block, m_row, m_col)`. `patch_keep_vector` reproduces that order.
+* A merged token whose `merge_size x merge_size` window straddles the mask boundary still mixes
+  unpainted patches, so about a token of contamination survives at the edge.
+"""
+
+import inspect
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+import numpy as np
+import torch
+
+from diffusers.utils import logging
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# A ContextVar rather than a global so a nested or concurrent encode cannot inherit another call's masks.
+_ACTIVE_KEEPS: ContextVar[list | None] = ContextVar("krea2_vision_keeps", default=None)
+
+_ORIGINAL_FORWARD = None
+_PATCHED = False
+
+
+def patch_keep_vector(keep_grid: np.ndarray, merge_size: int, temporal: int = 1) -> np.ndarray:
+    """Reorder a `(grid_h, grid_w)` keep grid into the processor's flattened patch order.
+
+    Sequence order is `(h_block, w_block, m_row, m_col)`, so a row-major flatten is wrong wherever
+    `merge_size > 1`.
+    """
+    grid_height, grid_width = keep_grid.shape
+    blocked = keep_grid.reshape(grid_height // merge_size, merge_size, grid_width // merge_size, merge_size)
+    ordered = blocked.transpose(0, 2, 1, 3).reshape(-1)
+    return np.tile(ordered, temporal) if temporal > 1 else ordered
+
+
+def build_patch_keeps(masks, image_grid_thw, merge_size: int, mask_to_grid) -> list:
+    """One patch keep vector per row of `image_grid_thw` (bool, processor patch order), or `None` per slot.
+
+    `mask_to_grid(mask, grid_height, grid_width)` rasterizes a mask onto a grid; it is injected so this
+    module carries no mask-format dependency.
+    """
+    if image_grid_thw is None:
+        return []
+    masks = list(masks or [])
+    keeps = []
+    for i in range(int(image_grid_thw.shape[0])):
+        mask = masks[i] if i < len(masks) else None
+        t, h, w = (int(x) for x in image_grid_thw[i].tolist())
+        grid = mask_to_grid(mask, h, w) if mask is not None else None
+        if grid is None:
+            keeps.append(None)
+            continue
+        keeps.append(torch.from_numpy(patch_keep_vector(grid, merge_size, t)).bool())
+    return keeps
+
+
+@contextmanager
+def vision_key_masks(keeps):
+    """Activate per-image patch key masks for the vision attention inside this block."""
+    if not keeps or all(keep is None for keep in keeps):
+        yield False
+        return
+    if not ensure_vision_attention_patched():
+        logger.warning(
+            "Qwen3-VL vision-attention patch unavailable; reference masks will only zero output tokens, "
+            "which cannot remove content from a reference."
+        )
+        yield False
+        return
+    token = _ACTIVE_KEEPS.set(list(keeps))
+    try:
+        yield True
+    finally:
+        _ACTIVE_KEEPS.reset(token)
+
+
+def _key_bias(keep: torch.Tensor, seq_len: int, device, dtype):
+    """Additive `(1, 1, 1, seq_len)` bias: 0 for kept keys, dtype-min for dropped ones.
+
+    Additive rather than a bool mask because the eager attention fallback adds the mask to the scores.
+    Only keys are masked, never queries, so no row can be fully masked and softmax cannot produce NaN.
+    """
+    if int(keep.numel()) != seq_len:
+        logger.warning(
+            f"Vision key mask length {int(keep.numel())} != patch sequence {seq_len}; skipping it for this image."
+        )
+        return None
+    bias = torch.zeros(1, 1, 1, seq_len, device=device, dtype=dtype)
+    bias.masked_fill_(~keep.to(device).view(1, 1, 1, seq_len), torch.finfo(dtype).min)
+    return bias
+
+
+def _forward_with_key_mask(self, hidden_states, cu_seqlens, position_embeddings=None, **kwargs):
+    keeps = _ACTIVE_KEEPS.get()
+    if not keeps:
+        return _ORIGINAL_FORWARD(
+            self,
+            hidden_states,
+            cu_seqlens,
+            position_embeddings=position_embeddings,
+            **kwargs,
+        )
+
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+    from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+        apply_rotary_pos_emb_vision,
+        eager_attention_forward,
+    )
+
+    seq_length = hidden_states.shape[0]
+    query_states, key_states, value_states = (
+        self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+    )
+    cos, sin = position_embeddings
+    query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+
+    query_states = query_states.transpose(0, 1).unsqueeze(0)
+    key_states = key_states.transpose(0, 1).unsqueeze(0)
+    value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+    attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+        self.config._attn_implementation, eager_attention_forward
+    )
+
+    # Always the per-image split path, never the flash `cu_seqlens` path: a key mask is per image, and
+    # varlen flash attention takes no mask. This is the same partition upstream uses for SDPA.
+    lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+    splits = [torch.split(tensor, lengths, dim=2) for tensor in (query_states, key_states, value_states)]
+
+    attn_outputs = []
+    for i, (q, k, v) in enumerate(zip(*splits)):
+        keep = keeps[i] if i < len(keeps) else None
+        bias = None if keep is None else _key_bias(keep, q.shape[2], q.device, q.dtype)
+        attn_outputs.append(
+            attention_interface(
+                self,
+                q,
+                k,
+                v,
+                attention_mask=bias,
+                scaling=self.scaling,
+                dropout=0.0,
+                is_causal=False,
+                **kwargs,
+            )[0]
+        )
+
+    attn_output = torch.cat(attn_outputs, dim=1)
+    attn_output = attn_output.reshape(seq_length, -1).contiguous()
+    return self.proj(attn_output)
+
+
+def ensure_vision_attention_patched() -> bool:
+    """Patch `Qwen3VLVisionAttention.forward` once. Returns False if the shape it relies on changed."""
+    global _ORIGINAL_FORWARD, _PATCHED
+    if _PATCHED:
+        return True
+    try:
+        from transformers.models.qwen3_vl.modeling_qwen3_vl import (
+            Qwen3VLVisionAttention,
+        )
+    except ImportError:
+        return False
+
+    # `_forward_with_key_mask` reproduces this signature and these attributes. If transformers changes
+    # either, skip patching and degrade to output-only masking rather than crash.
+    forward = getattr(Qwen3VLVisionAttention, "forward", None)
+    if forward is None:
+        return False
+    try:
+        params = inspect.signature(forward).parameters
+        init_src = inspect.getsource(Qwen3VLVisionAttention.__init__)
+    except (TypeError, ValueError, OSError):
+        return False
+    if not {"hidden_states", "cu_seqlens", "position_embeddings"}.issubset(params):
+        return False
+    if not all(f"self.{attr}" in init_src for attr in ("qkv", "proj", "num_heads", "scaling", "config")):
+        return False
+
+    _ORIGINAL_FORWARD = forward
+    Qwen3VLVisionAttention.forward = _forward_with_key_mask
+    _PATCHED = True
+    logger.debug("Patched Qwen3VLVisionAttention.forward for reference key masking.")
+    return True

modular_config.json

@@ -0,0 +1,10 @@
+{
+  "_class_name": "Krea2TurboReferenceAutoBlocks",
+  "_diffusers_version": "0.41.0.dev0",
+  "auto_map": {
+    "ModularPipelineBlocks": "krea2_reference.Krea2TurboReferenceAutoBlocks"
+  },
+  "requirements": {
+    "sdnq": ">=0.2.0"
+  }
+}

modular_model_index.json

@@ -0,0 +1,75 @@
+{
+  "_blocks_class_name": "Krea2TurboReferenceAutoBlocks",
+  "_class_name": "Krea2ModularPipeline",
+  "_diffusers_version": "0.41.0.dev0",
+  "scheduler": [
+    "diffusers",
+    "FlowMatchEulerDiscreteScheduler",
+    {
+      "pretrained_model_name_or_path": "OzzyGT/Krea_2_Turbo_sdnq_dynamic_8bit",
+      "revision": null,
+      "subfolder": "scheduler",
+      "type_hint": [
+        "diffusers",
+        "FlowMatchEulerDiscreteScheduler"
+      ],
+      "variant": null
+    }
+  ],
+  "text_encoder": [
+    "transformers",
+    "Qwen3VLModel",
+    {
+      "pretrained_model_name_or_path": "OzzyGT/Krea_2_Turbo_sdnq_dynamic_8bit",
+      "revision": null,
+      "subfolder": "text_encoder",
+      "type_hint": [
+        "transformers",
+        "Qwen3VLModel"
+      ],
+      "variant": null
+    }
+  ],
+  "tokenizer": [
+    "transformers",
+    "Qwen2Tokenizer",
+    {
+      "pretrained_model_name_or_path": "OzzyGT/Krea_2_Turbo_sdnq_dynamic_8bit",
+      "revision": null,
+      "subfolder": "tokenizer",
+      "type_hint": [
+        "transformers",
+        "Qwen2Tokenizer"
+      ],
+      "variant": null
+    }
+  ],
+  "transformer": [
+    "diffusers",
+    "Krea2Transformer2DModel",
+    {
+      "pretrained_model_name_or_path": "OzzyGT/Krea_2_Turbo_sdnq_dynamic_8bit",
+      "revision": null,
+      "subfolder": "transformer",
+      "type_hint": [
+        "diffusers",
+        "Krea2Transformer2DModel"
+      ],
+      "variant": null
+    }
+  ],
+  "vae": [
+    "diffusers",
+    "AutoencoderKLQwenImage",
+    {
+      "pretrained_model_name_or_path": "OzzyGT/Krea_2_Turbo_sdnq_dynamic_8bit",
+      "revision": null,
+      "subfolder": "vae",
+      "type_hint": [
+        "diffusers",
+        "AutoencoderKLQwenImage"
+      ],
+      "variant": null
+    }
+  ]
+}