krea2_reference.py

# 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."
        )