# Copyright 2026 Krea AI and 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.

"""
Text and VAE encoder blocks for Krea 2 pipelines.
"""

import math

import numpy as np
import PIL.Image
import torch
import torch.nn.functional as F
from transformers import Qwen2Tokenizer, Qwen3VLModel, Qwen3VLProcessor

from diffusers.configuration_utils import FrozenDict
from diffusers.guiders import ClassifierFreeGuidance
from diffusers.image_processor import InpaintProcessor, VaeImageProcessor
from diffusers.models import AutoencoderKLQwenImage
from diffusers.utils import logging
from diffusers.modular_pipelines.modular_pipeline import ModularPipelineBlocks, PipelineState
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
from .modular_pipeline import Krea2ModularPipeline, Krea2Pachifier


logger = logging.get_logger(__name__)

# Text conditioning uses the Qwen-Image chat template, tokenized as a fixed-length block: the prompt is padded to a
# fixed length first and the assistant suffix is appended after the padding (matching how the model was sampled at
# training time). The first `KREA2_PROMPT_TEMPLATE_START_IDX` (system prefix) tokens are dropped from the encoder
# outputs.
KREA2_PROMPT_TEMPLATE_PREFIX = (
    "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, "
    "spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n"
)
KREA2_PROMPT_TEMPLATE_SUFFIX = "<|im_end|>\n<|im_start|>assistant\n"
KREA2_PROMPT_TEMPLATE_START_IDX = 34
KREA2_PROMPT_TEMPLATE_NUM_SUFFIX_TOKENS = 5

# Indices into the text encoder's `hidden_states` tuple (0 is the embedding output) whose states are stacked per token
# and fed to the transformer's text fusion stage. These are the Krea 2 (Qwen3-VL-4B) taps; must have
# `transformer.config.num_text_layers` entries.
KREA2_TEXT_ENCODER_SELECT_LAYERS = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)


def get_krea2_prompt_embeds(
    text_encoder,
    tokenizer,
    prompt: str | list[str],
    text_encoder_select_layers: tuple[int, ...] = KREA2_TEXT_ENCODER_SELECT_LAYERS,
    prompt_template_prefix: str = KREA2_PROMPT_TEMPLATE_PREFIX,
    prompt_template_suffix: str = KREA2_PROMPT_TEMPLATE_SUFFIX,
    prompt_template_start_idx: int = KREA2_PROMPT_TEMPLATE_START_IDX,
    prompt_template_num_suffix_tokens: int = KREA2_PROMPT_TEMPLATE_NUM_SUFFIX_TOKENS,
    max_sequence_length: int = 512,
    device: torch.device | None = None,
):
    """Tokenize `prompt` into the fixed-length Krea 2 layout and tap the selected encoder hidden states.

    Returns a `(prompt_embeds, prompt_embeds_mask)` tuple of shapes `(batch_size, text_seq_len, num_text_layers,
    text_hidden_dim)` and `(batch_size, text_seq_len)` (bool).
    """
    prompt = [prompt] if isinstance(prompt, str) else prompt
    prefix_idx = prompt_template_start_idx
    text = [prompt_template_prefix + e for e in prompt]
    text_tokens = tokenizer(
        text,
        truncation=True,
        padding="max_length",
        max_length=max_sequence_length + prefix_idx - prompt_template_num_suffix_tokens,
        return_tensors="pt",
    ).to(device)
    suffix_tokens = tokenizer([prompt_template_suffix] * len(text), return_tensors="pt").to(device)

    input_ids = torch.cat([text_tokens.input_ids, suffix_tokens.input_ids], dim=1)
    attention_mask = torch.cat([text_tokens.attention_mask, suffix_tokens.attention_mask], dim=1).bool()

    # Krea 2 pads in the middle of the template (`[prefix | prompt | PAD | suffix]`), so the suffix tokens sit
    # downstream of the padding. The text features must use positions that count only real tokens (padding does
    # not consume a position) to match how the model was trained; otherwise the suffix gets a shifted mRoPE phase.
    # `Qwen3VLModel`'s default raw-index positions would place the suffix at ~max_length instead. Build the
    # cumulative-valid-token positions explicitly and broadcast across the 3 mRoPE axes (T/H/W are equal for text).
    position_ids = (attention_mask.long().cumsum(dim=-1) - 1).clamp(min=0)
    position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)

    outputs = text_encoder(
        input_ids=input_ids,
        attention_mask=attention_mask,
        position_ids=position_ids,
        output_hidden_states=True,
    )
    hidden_states = torch.stack([outputs.hidden_states[i] for i in text_encoder_select_layers], dim=2)

    prompt_embeds = hidden_states[:, prefix_idx:]
    prompt_embeds_mask = attention_mask[:, prefix_idx:]
    return prompt_embeds, prompt_embeds_mask


# Reference images ride in the user message ahead of the prompt through named vision placeholders; the Qwen3-VL
# processor expands each `<|image_pad|>` into the image's token grid so the text conditioning "sees" the references.
KREA2_EDIT_IMAGE_PLACEHOLDER = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"


def to_chw_tensor(image) -> torch.Tensor:
    """Convert a PIL image / numpy array / CHW tensor to a float CHW tensor in [0, 1]."""
    if isinstance(image, torch.Tensor):
        t = image.squeeze(0) if image.ndim == 4 else image
        t = t.float()
        if t.min() < 0:  # assume [-1, 1]
            t = (t + 1.0) / 2.0
        return t.clamp(0, 1)
    if isinstance(image, np.ndarray):
        image = PIL.Image.fromarray(image)
    image = image.convert("RGB")
    arr = np.asarray(image).astype(np.float32) / 255.0
    return torch.from_numpy(arr).permute(2, 0, 1)


def prep_vl_images(images: list[torch.Tensor], max_pixels: int) -> list[torch.Tensor]:
    """Resize reference images for the Qwen3-VL pass: aspect-preserving downscale (never upscaled) to fit
    `max_pixels` total area. The MLLM only needs a coarse view of the references; high-res detail flows through the
    VAE reference latents."""
    prepped = []
    for img in images:
        h, w = img.shape[1], img.shape[2]
        scale = min(1.0, math.sqrt(max_pixels / (h * w)))
        nh, nw = max(round(h * scale), 28), max(round(w * scale), 28)
        if (nh, nw) != (h, w):
            img = F.interpolate(img.unsqueeze(0).float(), size=(nh, nw), mode="bicubic", antialias=True).squeeze(0)
            img = img.clamp(0, 1)
        prepped.append(img.float())
    return prepped


def get_krea2_edit_prompt_embeds(
    text_encoder,
    tokenizer,
    processor,
    prompt: str | list[str],
    images: list[torch.Tensor] | None = None,
    text_encoder_select_layers: tuple[int, ...] = KREA2_TEXT_ENCODER_SELECT_LAYERS,
    prompt_template_prefix: str = KREA2_PROMPT_TEMPLATE_PREFIX,
    prompt_template_suffix: str = KREA2_PROMPT_TEMPLATE_SUFFIX,
    prompt_template_start_idx: int = KREA2_PROMPT_TEMPLATE_START_IDX,
    max_sequence_length: int = 512,
    device: torch.device | None = None,
):
    """Encode prompts for the edit task, embedding reference images (a coarse VL view of each) into the text
    conditioning through the Qwen3-VL vision tower.

    Unlike `get_krea2_prompt_embeds`, prompts are tokenized at their natural (unpadded) length: the processor
    expands each `<|image_pad|>` placeholder into a run of vision tokens that must stay intact, so no truncation or
    fixed padding is applied before encoding. All prompts share the same `images`. Returns a
    `(prompt_embeds, prompt_embeds_mask)` tuple of shapes `(batch_size, text_seq_len, num_text_layers,
    text_hidden_dim)` and `(batch_size, text_seq_len)` (bool), right-padded across the batch.
    """
    prompt = [prompt] if isinstance(prompt, str) else prompt
    prefix_idx = prompt_template_start_idx

    # The suffix is tokenized separately so it lands after the (image +) prompt tokens.
    suffix_inputs = tokenizer([prompt_template_suffix], return_tensors="pt").to(device)
    suffix_ids = suffix_inputs["input_ids"]
    suffix_mask = suffix_inputs["attention_mask"].bool()

    image_prompt = ""
    if images:
        image_prompt = "".join(KREA2_EDIT_IMAGE_PLACEHOLDER.format(i + 1) for i in range(len(images)))

    features = []
    for p in prompt:
        text = prompt_template_prefix + image_prompt + p
        extra_inputs = {}
        if images:
            # No truncation: the expanded image-pad runs must stay intact.
            inputs = processor(text=[text], images=list(images), return_tensors="pt", do_rescale=False).to(device)
            for k, v in inputs.items():
                if k in ("input_ids", "attention_mask"):
                    continue
                if isinstance(v, torch.Tensor) and v.is_floating_point():
                    v = v.to(text_encoder.dtype)
                extra_inputs[k] = v
        else:
            inputs = tokenizer(
                [text], truncation=True, max_length=max_sequence_length + prefix_idx, return_tensors="pt"
            ).to(device)

        input_ids = torch.cat([inputs["input_ids"], suffix_ids], dim=1)
        attention_mask = torch.cat([inputs["attention_mask"].bool(), suffix_mask], dim=1)

        # mm_token_type_ids (used for M-RoPE) must cover the appended suffix tokens too; they are plain text -> type 0.
        if "mm_token_type_ids" in extra_inputs:
            tt = extra_inputs["mm_token_type_ids"]
            extra_inputs["mm_token_type_ids"] = torch.cat([tt, torch.zeros_like(suffix_ids, dtype=tt.dtype)], dim=1)

        outputs = text_encoder(
            input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True, **extra_inputs
        )
        hidden_states = torch.stack([outputs.hidden_states[i] for i in text_encoder_select_layers], dim=2)
        # Drop the system-prefix tokens; what remains is (image +) prompt + suffix.
        features.append(hidden_states[0, prefix_idx:])

    max_len = max(f.shape[0] for f in features)
    prompt_embeds = features[0].new_zeros(len(features), max_len, *features[0].shape[1:])
    prompt_embeds_mask = torch.zeros(len(features), max_len, dtype=torch.bool, device=device)
    for i, f in enumerate(features):
        prompt_embeds[i, : f.shape[0]] = f
        prompt_embeds_mask[i, : f.shape[0]] = True

    return prompt_embeds, prompt_embeds_mask


# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
def retrieve_latents(
    encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample"
):
    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
        return encoder_output.latent_dist.sample(generator)
    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
        return encoder_output.latent_dist.mode()
    elif hasattr(encoder_output, "latents"):
        return encoder_output.latents
    else:
        raise AttributeError("Could not access latents of provided encoder_output")


# Modified from diffusers.modular_pipelines.qwenimage.encoders.encode_vae_image
def encode_vae_image(
    image: torch.Tensor,
    vae: AutoencoderKLQwenImage,
    generator: torch.Generator,
    device: torch.device,
    dtype: torch.dtype,
    latent_channels: int = 16,
    sample_mode: str = "argmax",
):
    if not isinstance(image, torch.Tensor):
        raise ValueError(f"Expected image to be a tensor, got {type(image)}.")

    # preprocessed image should be a 4D tensor: batch_size, num_channels, height, width
    if image.dim() == 4:
        image = image.unsqueeze(2)
    elif image.dim() != 5:
        raise ValueError(f"Expected image dims 4 or 5, got {image.dim()}.")

    image = image.to(device=device, dtype=dtype)

    if isinstance(generator, list):
        image_latents = [
            retrieve_latents(vae.encode(image[i : i + 1]), generator=generator[i], sample_mode=sample_mode)
            for i in range(image.shape[0])
        ]
        image_latents = torch.cat(image_latents, dim=0)
    else:
        image_latents = retrieve_latents(vae.encode(image), generator=generator, sample_mode=sample_mode)
    latents_mean = (
        torch.tensor(vae.config.latents_mean)
        .view(1, latent_channels, 1, 1, 1)
        .to(image_latents.device, image_latents.dtype)
    )
    latents_std = (
        torch.tensor(vae.config.latents_std)
        .view(1, latent_channels, 1, 1, 1)
        .to(image_latents.device, image_latents.dtype)
    )
    image_latents = (image_latents - latents_mean) / latents_std

    return image_latents


def encode_reference_latents(
    images: list[torch.Tensor],
    vae: AutoencoderKLQwenImage,
    max_pixels: int,
    generator: torch.Generator | None,
    device: torch.device,
    vae_scale_factor: int,
    patch_size: int,
    latent_channels: int = 16,
) -> list[torch.Tensor]:
    """Encode `[0, 1]` CHW reference images to normalized VAE latents, one `(C, h, w)` tensor per image. Each image
    is downscaled (aspect-preserving, never upscaled) to fit within `max_pixels`, then snapped so the latent grid is
    patchifiable. References keep their own aspect ratio, independent of the generated output size."""
    snap = vae_scale_factor * patch_size
    vae_dtype = vae.dtype

    latents_mean = torch.tensor(vae.config.latents_mean).view(1, latent_channels, 1, 1, 1)
    latents_std = torch.tensor(vae.config.latents_std).view(1, latent_channels, 1, 1, 1)

    ref_latents = []
    for img in images:
        img = img.unsqueeze(0).to(device, dtype=vae_dtype)
        h, w = img.shape[2], img.shape[3]
        if h * w > max_pixels:
            ratio = h / w
            new_h, new_w = math.sqrt(max_pixels * ratio), math.sqrt(max_pixels / ratio)
        else:
            new_h, new_w = float(h), float(w)
        new_h = max(snap, int(round(new_h / snap)) * snap)
        new_w = max(snap, int(round(new_w / snap)) * snap)
        if (new_h, new_w) != (h, w):
            img = F.interpolate(img.float(), size=(new_h, new_w), mode="bilinear").to(vae_dtype)

        img = (img * 2.0 - 1.0).unsqueeze(2)  # [0, 1] -> [-1, 1], add frame dim
        latent = retrieve_latents(vae.encode(img), generator=generator, sample_mode="sample")
        latent = (latent - latents_mean.to(latent.device, latent.dtype)) / latents_std.to(latent.device, latent.dtype)
        ref_latents.append(latent[:, :, 0][0])  # drop frame + batch dims -> (C, h, w)
    return ref_latents


def pack_reference_latents(
    ref_latents: list[torch.Tensor],
    pachifier: Krea2Pachifier,
    device: torch.device,
    dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Patchify reference latents into `(1, ref_seq_len, C * p * p)` tokens and build their `(ref_seq_len, 3)` rotary
    coordinates. The i-th reference sits on frame axis `i + 1` with its own y/x grid starting at 0 (the Kontext-style
    "index" placement that marks each reference as a distinct image rather than more of the canvas)."""
    p = pachifier.config.patch_size
    tokens, position_ids = [], []
    for i, ref in enumerate(ref_latents):
        ref = ref.unsqueeze(0).to(device, dtype)
        tokens.append(pachifier.pack_latents(ref))
        _, _, h, w = ref.shape
        ids = torch.zeros(h // p, w // p, 3, device=device)
        ids[..., 0] = i + 1
        ids[..., 1] = torch.arange(h // p, device=device)[:, None]
        ids[..., 2] = torch.arange(w // p, device=device)[None, :]
        position_ids.append(ids.reshape(-1, 3))
    return torch.cat(tokens, dim=1), torch.cat(position_ids, dim=0)


# ====================
# 1. TEXT ENCODER
# ====================


class Krea2TextEncoderStep(ModularPipelineBlocks):
    model_name = "krea2"

    def __init__(self, text_encoder_select_layers: tuple[int, ...] | None = None):
        """Text encoder step for Krea 2.

        Args:
            text_encoder_select_layers (`tuple[int, ...]`, *optional*):
                Indices into the text encoder's `hidden_states` tuple (0 is the embedding output) whose states are
                stacked per token as the transformer's text conditioning. Must have
                `transformer.config.num_text_layers` entries. Defaults to the Krea 2 (Qwen3-VL-4B) taps.
        """
        if text_encoder_select_layers is None:
            text_encoder_select_layers = KREA2_TEXT_ENCODER_SELECT_LAYERS
        self.text_encoder_select_layers = tuple(text_encoder_select_layers)
        super().__init__()

    @property
    def description(self) -> str:
        return "Text Encoder step that generates text embeddings to guide the image generation."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen3VLModel, description="The text encoder to use"),
            ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("prompt"),
            InputParam.template("negative_prompt"),
            InputParam.template("max_sequence_length"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam.template("prompt_embeds"),
            OutputParam.template("prompt_embeds_mask"),
            OutputParam.template("negative_prompt_embeds"),
            OutputParam.template("negative_prompt_embeds_mask"),
        ]

    @staticmethod
    def check_inputs(prompt, negative_prompt, max_sequence_length):
        if not isinstance(prompt, str) and not isinstance(prompt, list):
            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")

        if (
            negative_prompt is not None
            and not isinstance(negative_prompt, str)
            and not isinstance(negative_prompt, list)
        ):
            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")

        if max_sequence_length is not None and max_sequence_length <= 0:
            raise ValueError(f"`max_sequence_length` must be a positive integer but is {max_sequence_length}")

    @torch.no_grad()
    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        device = components._execution_device
        self.check_inputs(block_state.prompt, block_state.negative_prompt, block_state.max_sequence_length)

        block_state.prompt_embeds, block_state.prompt_embeds_mask = get_krea2_prompt_embeds(
            components.text_encoder,
            components.tokenizer,
            prompt=block_state.prompt,
            text_encoder_select_layers=self.text_encoder_select_layers,
            max_sequence_length=block_state.max_sequence_length,
            device=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 or ""
            block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_krea2_prompt_embeds(
                components.text_encoder,
                components.tokenizer,
                prompt=negative_prompt,
                text_encoder_select_layers=self.text_encoder_select_layers,
                max_sequence_length=block_state.max_sequence_length,
                device=device,
            )

        self.set_block_state(state, block_state)
        return components, state


class Krea2EditTextEncoderStep(ModularPipelineBlocks):
    model_name = "krea2"

    def __init__(self, text_encoder_select_layers: tuple[int, ...] | None = None):
        """Text encoder step for the Krea 2 edit task: encodes the prompt while embedding a coarse view of the
        reference image(s) into the conditioning through the Qwen3-VL vision tower.

        Args:
            text_encoder_select_layers (`tuple[int, ...]`, *optional*):
                Indices into the text encoder's `hidden_states` tuple whose states are stacked per token as the
                transformer's text conditioning. Defaults to the Krea 2 (Qwen3-VL-4B) taps.
        """
        if text_encoder_select_layers is None:
            text_encoder_select_layers = KREA2_TEXT_ENCODER_SELECT_LAYERS
        self.text_encoder_select_layers = tuple(text_encoder_select_layers)
        super().__init__()

    @property
    def description(self) -> str:
        return (
            "Text encoder step for the edit task. Embeds reference image(s) into the text conditioning via the "
            "Qwen3-VL vision tower, matching how the Ostris AI-Toolkit edit LoRAs are trained."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("text_encoder", Qwen3VLModel, description="The text encoder to use"),
            ComponentSpec("tokenizer", Qwen2Tokenizer, description="The tokenizer to use"),
            ComponentSpec("processor", Qwen3VLProcessor, description="The Qwen3-VL processor for reference images"),
            ComponentSpec(
                "guider",
                ClassifierFreeGuidance,
                config=FrozenDict({"guidance_scale": 4.5, "use_original_formulation": True}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("prompt"),
            InputParam.template("negative_prompt"),
            InputParam.template("image", required=True, note="The reference image(s) for the edit."),
            InputParam.template("max_sequence_length"),
            InputParam(
                "vl_image_max_pixels",
                type_hint=int,
                default=384 * 384,
                description="Pixel budget for the coarse Qwen3-VL view of each reference image.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam.template("prompt_embeds"),
            OutputParam.template("prompt_embeds_mask"),
            OutputParam.template("negative_prompt_embeds"),
            OutputParam.template("negative_prompt_embeds_mask"),
        ]

    @staticmethod
    def check_inputs(prompt, negative_prompt, max_sequence_length):
        if not isinstance(prompt, str) and not isinstance(prompt, list):
            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
        if (
            negative_prompt is not None
            and not isinstance(negative_prompt, str)
            and not isinstance(negative_prompt, list)
        ):
            raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}")
        if max_sequence_length is not None and max_sequence_length <= 0:
            raise ValueError(f"`max_sequence_length` must be a positive integer but is {max_sequence_length}")

    @torch.no_grad()
    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        device = components._execution_device
        self.check_inputs(block_state.prompt, block_state.negative_prompt, block_state.max_sequence_length)

        image_list = block_state.image if isinstance(block_state.image, (list, tuple)) else [block_state.image]
        ref_images = [to_chw_tensor(img).to(device) for img in image_list]
        vl_images = prep_vl_images(ref_images, block_state.vl_image_max_pixels)

        block_state.prompt_embeds, block_state.prompt_embeds_mask = get_krea2_edit_prompt_embeds(
            components.text_encoder,
            components.tokenizer,
            components.processor,
            prompt=block_state.prompt,
            images=vl_images,
            text_encoder_select_layers=self.text_encoder_select_layers,
            max_sequence_length=block_state.max_sequence_length,
            device=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 or ""
            block_state.negative_prompt_embeds, block_state.negative_prompt_embeds_mask = get_krea2_edit_prompt_embeds(
                components.text_encoder,
                components.tokenizer,
                components.processor,
                prompt=negative_prompt,
                images=vl_images,
                text_encoder_select_layers=self.text_encoder_select_layers,
                max_sequence_length=block_state.max_sequence_length,
                device=device,
            )

        self.set_block_state(state, block_state)
        return components, state


# ====================
# 2. IMAGE PREPROCESS
# ====================


class Krea2InpaintProcessImagesInputStep(ModularPipelineBlocks):
    model_name = "krea2"

    @property
    def description(self) -> str:
        return "Image Preprocess step for inpainting task. This processes the image and mask inputs together. Images will be resized to the given height and width."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_mask_processor",
                InpaintProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("mask_image"),
            InputParam.template("image"),
            InputParam.template("height"),
            InputParam.template("width"),
            InputParam.template("padding_mask_crop"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            ),
            OutputParam(
                name="processed_mask_image",
                type_hint=torch.Tensor,
                description="The processed mask image",
            ),
            OutputParam(
                name="mask_overlay_kwargs",
                type_hint=dict,
                description="The kwargs for the postprocess step to apply the mask overlay",
            ),
        ]

    @staticmethod
    def check_inputs(height, width, vae_scale_factor):
        if height is not None and height % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")

        if width is not None and width % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")

    @torch.no_grad()
    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        self.check_inputs(
            height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
        )
        height = block_state.height or components.default_height
        width = block_state.width or components.default_width

        block_state.processed_image, block_state.processed_mask_image, block_state.mask_overlay_kwargs = (
            components.image_mask_processor.preprocess(
                image=block_state.image,
                mask=block_state.mask_image,
                height=height,
                width=width,
                padding_mask_crop=block_state.padding_mask_crop,
            )
        )

        self.set_block_state(state, block_state)
        return components, state


class Krea2ProcessImagesInputStep(ModularPipelineBlocks):
    model_name = "krea2"

    @property
    def description(self) -> str:
        return "Image Preprocess step. will resize the image to the given height and width."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec(
                "image_processor",
                VaeImageProcessor,
                config=FrozenDict({"vae_scale_factor": 16}),
                default_creation_method="from_config",
            ),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("image"),
            InputParam.template("height"),
            InputParam.template("width"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="processed_image",
                type_hint=torch.Tensor,
                description="The processed image",
            )
        ]

    @staticmethod
    def check_inputs(height, width, vae_scale_factor):
        if height is not None and height % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Height must be divisible by {vae_scale_factor * 2} but is {height}")

        if width is not None and width % (vae_scale_factor * 2) != 0:
            raise ValueError(f"Width must be divisible by {vae_scale_factor * 2} but is {width}")

    @torch.no_grad()
    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
        block_state = self.get_block_state(state)

        self.check_inputs(
            height=block_state.height, width=block_state.width, vae_scale_factor=components.vae_scale_factor
        )
        height = block_state.height or components.default_height
        width = block_state.width or components.default_width

        block_state.processed_image = components.image_processor.preprocess(
            image=block_state.image,
            height=height,
            width=width,
        )

        self.set_block_state(state, block_state)
        return components, state


# ====================
# 3. VAE ENCODER
# ====================


class Krea2VaeEncoderStep(ModularPipelineBlocks):
    model_name = "krea2"

    @property
    def description(self) -> str:
        return "VAE Encoder step that converts processed_image into latent representations image_latents."

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [ComponentSpec("vae", AutoencoderKLQwenImage)]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam(
                name="processed_image", required=True, type_hint=torch.Tensor, description="The image tensor to encode"
            ),
            InputParam.template("generator"),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [OutputParam.template("image_latents")]

    @torch.no_grad()
    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        device = components._execution_device
        dtype = components.vae.dtype

        block_state.image_latents = encode_vae_image(
            image=block_state.processed_image,
            vae=components.vae,
            generator=block_state.generator,
            device=device,
            dtype=dtype,
            latent_channels=components.num_channels_latents,
        )

        self.set_block_state(state, block_state)

        return components, state


class Krea2EditReferenceLatentsStep(ModularPipelineBlocks):
    model_name = "krea2"

    @property
    def description(self) -> str:
        return (
            "Reference (edit) VAE encoder step. Encodes reference image(s) to clean, normalized VAE latents and packs "
            "them into transformer tokens with their frame-axis rotary coordinates. These tokens are appended to the "
            "sequence at flow time t=0 to condition the generation on the references."
        )

    @property
    def expected_components(self) -> list[ComponentSpec]:
        return [
            ComponentSpec("vae", AutoencoderKLQwenImage),
            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
        ]

    @property
    def inputs(self) -> list[InputParam]:
        return [
            InputParam.template("image", required=True, note="The reference image(s) for the edit."),
            InputParam.template("generator"),
            InputParam(
                "reference_max_pixels",
                type_hint=int,
                default=1024 * 1024,
                description="Pixel budget each reference image is downscaled to fit before VAE encoding.",
            ),
        ]

    @property
    def intermediate_outputs(self) -> list[OutputParam]:
        return [
            OutputParam(
                name="reference_latents",
                type_hint=torch.Tensor,
                description="Packed clean reference tokens of shape (1, ref_seq_len, C * p * p), appended to the "
                "denoiser sequence at t=0.",
            ),
            OutputParam(
                name="reference_position_ids",
                type_hint=torch.Tensor,
                description="Rotary coordinates (ref_seq_len, 3) for the reference tokens; the i-th reference sits on "
                "frame axis i + 1.",
            ),
            OutputParam(
                name="ref_seq_len",
                kwargs_type="denoiser_input_fields",
                type_hint=int,
                description="Number of reference tokens appended to the denoiser sequence.",
            ),
        ]

    @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.pachifier.config.patch_size

        image_list = block_state.image if isinstance(block_state.image, (list, tuple)) else [block_state.image]
        ref_images = [to_chw_tensor(img) for img in image_list]

        ref_latents = encode_reference_latents(
            images=ref_images,
            vae=components.vae,
            max_pixels=block_state.reference_max_pixels,
            generator=block_state.generator,
            device=device,
            vae_scale_factor=components.vae_scale_factor,
            patch_size=patch_size,
            latent_channels=components.num_channels_latents,
        )
        block_state.reference_latents, block_state.reference_position_ids = pack_reference_latents(
            ref_latents, components.pachifier, device, components.vae.dtype
        )
        block_state.ref_seq_len = block_state.reference_latents.shape[1]

        self.set_block_state(state, block_state)

        return components, state