Krea 2 edit modular blocks (remote code)

README.md

@@ -0,0 +1,62 @@
+---
+library_name: diffusers
+tags:
+  - modular-diffusers
+  - krea
+  - image-to-image
+  - image-editing
+base_model: krea/Krea-2-Turbo
+---
+
+# Krea 2 reference-image edit — Modular Diffusers blocks
+
+Custom [Modular Diffusers](https://huggingface.co/docs/diffusers/main/en/modular_diffusers/overview)
+blocks that reproduce the [`ostris/Krea2OstrisEdit`](https://huggingface.co/ostris/Krea2OstrisEdit)
+reference-image ("edit") workflow for **Krea 2**, loadable as remote code on top of stock `diffusers`.
+
+```python
+import torch
+from transformers import Qwen3VLProcessor
+from diffusers import ClassifierFreeGuidance
+from diffusers.modular_pipelines import ModularPipelineBlocks
+
+blocks = ModularPipelineBlocks.from_pretrained("diffusers-modular/krea2-edit", trust_remote_code=True)
+pipe = blocks.init_pipeline("krea/Krea-2-Turbo")        # weights from the base repo
+pipe.load_components(torch_dtype=torch.bfloat16)
+pipe.update_components(processor=Qwen3VLProcessor.from_pretrained("Qwen/Qwen3-VL-4B-Instruct"))
+pipe.update_components(guider=ClassifierFreeGuidance(guidance_scale=0.0, use_original_formulation=True))
+pipe.to("cuda")
+
+from PIL import Image
+image = pipe(
+    prompt="a white yeti with horns reading a book",
+    image=Image.open("reference.png"),   # one or more reference images
+    num_inference_steps=8, mu=1.15, output="images",
+)[0]
+```
+
+## What it does
+
+Reference images condition generation two ways (matching how the Ostris AI-Toolkit edit LoRAs train):
+
+1. **Qwen3-VL prompt embedding** — a coarse view of each reference is embedded into the text
+   conditioning through the vision tower.
+2. **Clean VAE latents at flow time t=0** — each reference is VAE-encoded and appended to the
+   transformer sequence as clean tokens on its own rotary frame axis, so the noisy image tokens
+   attend to it at every block.
+
+The bundled `Krea2Transformer2DModel` adds a small, backward-compatible `ref_seq_len` argument to the
+Krea 2 transformer forward (t=0 modulation of the reference span; those tokens are excluded from the
+predicted velocity). With `ref_seq_len=0` it is numerically identical to plain Krea 2 text-to-image.
+
+## Files
+
+- `block.py` — entry point (`Krea2EditBlocks`), referenced by `config.json`'s `auto_map`.
+- `transformer_krea2.py` — the Krea 2 transformer (with the `ref_seq_len` edit path).
+- `modular_blocks_krea2*.py`, `encoders.py`, `before_denoise.py`, `denoise.py`, `decoders.py`,
+  `inputs.py`, `modular_pipeline.py` — the modular blocks.
+
+Part of the [`diffusers-modular`](https://huggingface.co/diffusers-modular) effort to replicate
+popular ComfyUI / community workflows with Modular Diffusers.
+
+> Requires access to the gated `krea/Krea-2-Turbo` weights.

before_denoise.py

@@ -0,0 +1,635 @@
+# 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.
+
+import inspect
+
+import numpy as np
+import torch
+
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils.torch_utils import randn_tensor
+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
+
+
+# Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: int | None = None,
+    device: str | torch.device | None = None,
+    timesteps: list[int] | None = None,
+    sigmas: list[float] | None = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`list[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`list[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.before_denoise.get_timesteps
+def get_timesteps(scheduler, num_inference_steps, strength):
+    # get the original timestep using init_timestep
+    init_timestep = min(num_inference_steps * strength, num_inference_steps)
+
+    t_start = int(max(num_inference_steps - init_timestep, 0))
+    timesteps = scheduler.timesteps[t_start * scheduler.order :]
+    if hasattr(scheduler, "set_begin_index"):
+        scheduler.set_begin_index(t_start * scheduler.order)
+
+    return timesteps, num_inference_steps - t_start
+
+
+# ====================
+# 1. PREPARE LATENTS
+# ====================
+
+
+class Krea2PrepareLatentsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Prepare initial random noise for the generation process"
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("latents"),
+            InputParam.template("height"),
+            InputParam.template("width"),
+            InputParam.template("num_images_per_prompt"),
+            InputParam.template("generator"),
+            InputParam.template("batch_size"),
+            InputParam.template("dtype"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(name="height", type_hint=int, description="if not set, updated to default value"),
+            OutputParam(name="width", type_hint=int, description="if not set, updated to default value"),
+            OutputParam(
+                name="latents",
+                type_hint=torch.Tensor,
+                description="The initial latents to use for the denoising process",
+            ),
+        ]
+
+    @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) -> 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,
+        )
+
+        device = components._execution_device
+        batch_size = block_state.batch_size * block_state.num_images_per_prompt
+
+        # we can update the height and width here since it's used to generate the initial
+        block_state.height = block_state.height or components.default_height
+        block_state.width = block_state.width or components.default_width
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        latent_height = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+        latent_width = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+
+        shape = (batch_size, components.num_channels_latents, 1, latent_height, latent_width)
+        if isinstance(block_state.generator, list) and len(block_state.generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(block_state.generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        if block_state.latents is None:
+            block_state.latents = randn_tensor(
+                shape, generator=block_state.generator, device=device, dtype=block_state.dtype
+            )
+            block_state.latents = components.pachifier.pack_latents(block_state.latents)
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class Krea2PrepareLatentsWithStrengthStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that adds noise to image latents for image-to-image/inpainting. Should be run after set_timesteps, prepare_latents. Both noise and image latents should already be patchified."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial random noised, can be generated in prepare latent step.",
+            ),
+            InputParam.template("image_latents", note="Can be generated from vae encoder and updated in input step."),
+            InputParam(
+                name="timesteps",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="initial_noise",
+                type_hint=torch.Tensor,
+                description="The initial random noised used for inpainting denoising.",
+            ),
+            OutputParam(
+                name="latents",
+                type_hint=torch.Tensor,
+                description="The scaled noisy latents to use for inpainting/image-to-image denoising.",
+            ),
+        ]
+
+    @staticmethod
+    def check_inputs(image_latents, latents):
+        if image_latents.shape[0] != latents.shape[0]:
+            raise ValueError(
+                f"`image_latents` must have have same batch size as `latents`, but got {image_latents.shape[0]} and {latents.shape[0]}"
+            )
+
+        if image_latents.ndim != 3:
+            raise ValueError(f"`image_latents` must have 3 dimensions (patchified), but got {image_latents.ndim}")
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        self.check_inputs(
+            image_latents=block_state.image_latents,
+            latents=block_state.latents,
+        )
+
+        # prepare latent timestep
+        latent_timestep = block_state.timesteps[:1].repeat(block_state.latents.shape[0])
+
+        # make copy of initial_noise
+        block_state.initial_noise = block_state.latents
+
+        # scale noise
+        block_state.latents = components.scheduler.scale_noise(
+            block_state.image_latents, latent_timestep, block_state.latents
+        )
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class Krea2CreateMaskLatentsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that creates mask latents from preprocessed mask_image by interpolating to latent space."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="processed_mask_image",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The processed mask to use for the inpainting process.",
+            ),
+            InputParam.template("height", required=True),
+            InputParam.template("width", required=True),
+            InputParam.template("dtype"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="mask", type_hint=torch.Tensor, description="The mask to use for the inpainting process."
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        device = components._execution_device
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+
+        height_latents = 2 * (int(block_state.height) // (components.vae_scale_factor * 2))
+        width_latents = 2 * (int(block_state.width) // (components.vae_scale_factor * 2))
+
+        block_state.mask = torch.nn.functional.interpolate(
+            block_state.processed_mask_image,
+            size=(height_latents, width_latents),
+        )
+
+        block_state.mask = block_state.mask.unsqueeze(2)
+        block_state.mask = block_state.mask.repeat(1, components.num_channels_latents, 1, 1, 1)
+        block_state.mask = block_state.mask.to(device=device, dtype=block_state.dtype)
+
+        block_state.mask = components.pachifier.pack_latents(block_state.mask)
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+# ====================
+# 2. SET TIMESTEPS
+# ====================
+
+
+class Krea2SetTimestepsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that sets the scheduler's timesteps for text-to-image generation. Should be run after prepare latents step."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_inference_steps", default=28),
+            InputParam.template("sigmas"),
+            InputParam(
+                name="mu",
+                type_hint=float,
+                description="Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if not provided, computed from the image sequence length (base checkpoint behavior).",
+            ),
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial random noised latents for the denoising process. Can be generated in prepare latents step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="timesteps", type_hint=torch.Tensor, description="The timesteps to use for the denoising process"
+            ),
+        ]
+
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        device = components._execution_device
+        sigmas = (
+            np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+            if block_state.sigmas is None
+            else block_state.sigmas
+        )
+
+        mu = block_state.mu
+        if mu is None:
+            mu = calculate_shift(
+                image_seq_len=block_state.latents.shape[1],
+                base_seq_len=components.scheduler.config.get("base_image_seq_len", 256),
+                max_seq_len=components.scheduler.config.get("max_image_seq_len", 6400),
+                base_shift=components.scheduler.config.get("base_shift", 0.5),
+                max_shift=components.scheduler.config.get("max_shift", 1.15),
+            )
+        block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+            scheduler=components.scheduler,
+            num_inference_steps=block_state.num_inference_steps,
+            device=device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+
+        components.scheduler.set_begin_index(0)
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class Krea2SetTimestepsWithStrengthStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that sets the scheduler's timesteps for image-to-image generation, and inpainting. Should be run after prepare latents step."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_inference_steps", default=28),
+            InputParam.template("sigmas"),
+            InputParam(
+                name="mu",
+                type_hint=float,
+                description="Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if not provided, computed from the image sequence length (base checkpoint behavior).",
+            ),
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The latents to use for the denoising process. Can be generated in prepare latents step.",
+            ),
+            InputParam.template("strength"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="timesteps",
+                type_hint=torch.Tensor,
+                description="The timesteps to use for the denoising process.",
+            ),
+            OutputParam(
+                name="num_inference_steps",
+                type_hint=int,
+                description="The number of denoising steps to perform at inference time. Updated based on strength.",
+            ),
+        ]
+
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        device = components._execution_device
+        sigmas = (
+            np.linspace(1.0, 1 / block_state.num_inference_steps, block_state.num_inference_steps)
+            if block_state.sigmas is None
+            else block_state.sigmas
+        )
+
+        mu = block_state.mu
+        if mu is None:
+            mu = calculate_shift(
+                image_seq_len=block_state.latents.shape[1],
+                base_seq_len=components.scheduler.config.get("base_image_seq_len", 256),
+                max_seq_len=components.scheduler.config.get("max_image_seq_len", 6400),
+                base_shift=components.scheduler.config.get("base_shift", 0.5),
+                max_shift=components.scheduler.config.get("max_shift", 1.15),
+            )
+        block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps(
+            scheduler=components.scheduler,
+            num_inference_steps=block_state.num_inference_steps,
+            device=device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+
+        block_state.timesteps, block_state.num_inference_steps = get_timesteps(
+            scheduler=components.scheduler,
+            num_inference_steps=block_state.num_inference_steps,
+            strength=block_state.strength,
+        )
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+# ====================
+# 3. OTHER INPUTS FOR DENOISER
+# ====================
+
+## RoPE inputs for denoiser
+
+
+class Krea2RoPEInputsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step that prepares the rotary position ids for the denoising process. Text tokens sit at the origin, "
+            "image tokens carry their `(0, h, w)` latent-grid coordinates. Should be placed after prepare_latents step."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("height", required=True),
+            InputParam.template("width", required=True),
+            InputParam.template("prompt_embeds_mask"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="position_ids",
+                kwargs_type="denoiser_input_fields",
+                type_hint=torch.Tensor,
+                description="The rotary coordinates of shape (text_seq_len + grid_height * grid_width, 3) for the combined text-image sequence.",
+            ),
+        ]
+
+    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
+
+        text_seq_len = block_state.prompt_embeds_mask.shape[1]
+        grid_height = block_state.height // (components.vae_scale_factor * patch_size)
+        grid_width = block_state.width // (components.vae_scale_factor * patch_size)
+
+        text_ids = torch.zeros(text_seq_len, 3, device=device)
+        image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
+        image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
+        image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
+        image_ids = image_ids.reshape(grid_height * grid_width, 3)
+        block_state.position_ids = torch.cat([text_ids, image_ids], dim=0)
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class Krea2EditRoPEInputsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Step that prepares the rotary position ids for the edit task: the `[text, image]` coordinates followed "
+            "by the reference tokens' coordinates (each reference on its own frame axis). Should be placed after "
+            "prepare_latents and the reference latents step."
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("height", required=True),
+            InputParam.template("width", required=True),
+            InputParam.template("prompt_embeds_mask"),
+            InputParam(
+                "reference_position_ids",
+                required=True,
+                type_hint=torch.Tensor,
+                description="Rotary coordinates for the reference tokens. Can be generated in the reference latents step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="position_ids",
+                kwargs_type="denoiser_input_fields",
+                type_hint=torch.Tensor,
+                description="The rotary coordinates of shape (text_seq_len + grid_height * grid_width + ref_seq_len, 3) "
+                "for the combined text-image-reference sequence.",
+            ),
+        ]
+
+    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
+
+        text_seq_len = block_state.prompt_embeds_mask.shape[1]
+        grid_height = block_state.height // (components.vae_scale_factor * patch_size)
+        grid_width = block_state.width // (components.vae_scale_factor * patch_size)
+
+        text_ids = torch.zeros(text_seq_len, 3, device=device)
+        image_ids = torch.zeros(grid_height, grid_width, 3, device=device)
+        image_ids[..., 1] = torch.arange(grid_height, device=device)[:, None]
+        image_ids[..., 2] = torch.arange(grid_width, device=device)[None, :]
+        image_ids = image_ids.reshape(grid_height * grid_width, 3)
+        block_state.position_ids = torch.cat([text_ids, image_ids, block_state.reference_position_ids], dim=0)
+
+        self.set_block_state(state, block_state)
+
+        return components, state

block.py

@@ -0,0 +1,12 @@
+"""Entry point for the Krea 2 reference-image edit modular blocks (loaded via `trust_remote_code`).
+
+Importing this module also registers `Krea2Transformer2DModel` into the `diffusers` namespace so that
+`ModularPipeline` component loading can resolve the transformer referenced as
+`["diffusers", "Krea2Transformer2DModel"]` in the base `krea/Krea-2-Turbo` `model_index.json`.
+"""
+
+from .transformer_krea2 import Krea2Transformer2DModel  # noqa: F401  (registers into diffusers namespace)
+from .modular_blocks_krea2_edit import Krea2EditBlocks
+
+
+__all__ = ["Krea2EditBlocks", "Krea2Transformer2DModel"]

config.json

@@ -0,0 +1,5 @@
+{
+  "auto_map": {
+    "ModularPipelineBlocks": "block.Krea2EditBlocks"
+  }
+}

decoders.py

@@ -0,0 +1,265 @@
+# 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.
+
+
+from typing import Any
+
+import torch
+
+from diffusers.configuration_utils import FrozenDict
+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__)
+
+
+# after denoising loop (unpack latents)
+
+
+class Krea2AfterDenoiseStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that unpack the latents from 3D tensor (batch_size, sequence_length, channels) into 5D tensor (batch_size, channels, 1, height, width)"
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        components = [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+        return components
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("height", required=True),
+            InputParam.template("width", required=True),
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The latents to decode, can be generated in the denoise step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                name="latents", type_hint=torch.Tensor, description="The denoised latents unpacked to B, C, 1, H, W"
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        vae_scale_factor = components.vae_scale_factor
+        block_state.latents = components.pachifier.unpack_latents(
+            block_state.latents, block_state.height, block_state.width, vae_scale_factor=vae_scale_factor
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# decode step
+
+
+class Krea2DecoderStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "Step that decodes the latents to images"
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        components = [
+            ComponentSpec("vae", AutoencoderKLQwenImage),
+        ]
+
+        return components
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam.template("images", note="tensor output of the vae decoder.")]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        if block_state.latents.ndim == 4:
+            block_state.latents = block_state.latents.unsqueeze(dim=1)
+        elif block_state.latents.ndim != 5:
+            raise ValueError(
+                f"expect latents to be a 4D or 5D tensor but got: {block_state.latents.shape}. Please make sure the latents are unpacked before decode step."
+            )
+        block_state.latents = block_state.latents.to(components.vae.dtype)
+
+        latents_mean = (
+            torch.tensor(components.vae.config.latents_mean)
+            .view(1, components.vae.config.z_dim, 1, 1, 1)
+            .to(block_state.latents.device, block_state.latents.dtype)
+        )
+        latents_std = 1.0 / torch.tensor(components.vae.config.latents_std).view(
+            1, components.vae.config.z_dim, 1, 1, 1
+        ).to(block_state.latents.device, block_state.latents.dtype)
+        block_state.latents = block_state.latents / latents_std + latents_mean
+        block_state.images = components.vae.decode(block_state.latents, return_dict=False)[0][:, :, 0]
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+# postprocess the decoded images
+
+
+class Krea2ProcessImagesOutputStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "postprocess the generated image"
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "image_processor",
+                VaeImageProcessor,
+                config=FrozenDict({"vae_scale_factor": 16}),
+                default_creation_method="from_config",
+            ),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="images",
+                required=True,
+                type_hint=torch.Tensor,
+                description="the generated image tensor from decoders step",
+            ),
+            InputParam.template("output_type"),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam.template("images")]
+
+    @staticmethod
+    def check_inputs(output_type):
+        if output_type not in ["pil", "np", "pt"]:
+            raise ValueError(f"Invalid output_type: {output_type}")
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
+        block_state = self.get_block_state(state)
+
+        self.check_inputs(block_state.output_type)
+
+        block_state.images = components.image_processor.postprocess(
+            image=block_state.images,
+            output_type=block_state.output_type,
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state
+
+
+class Krea2InpaintProcessImagesOutputStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return "postprocess the generated image, optionally apply the mask overlay to the original image."
+
+    @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(
+                name="images",
+                required=True,
+                type_hint=torch.Tensor,
+                description="the generated image tensor from decoders step",
+            ),
+            InputParam.template("output_type"),
+            InputParam(
+                name="mask_overlay_kwargs",
+                type_hint=dict[str, Any],
+                description="The kwargs for the postprocess step to apply the mask overlay. generated in Krea2InpaintProcessImagesInputStep.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [OutputParam.template("images")]
+
+    @staticmethod
+    def check_inputs(output_type, mask_overlay_kwargs):
+        if output_type not in ["pil", "np", "pt"]:
+            raise ValueError(f"Invalid output_type: {output_type}")
+
+        if mask_overlay_kwargs and output_type != "pil":
+            raise ValueError("only support output_type 'pil' for mask overlay")
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState):
+        block_state = self.get_block_state(state)
+
+        self.check_inputs(block_state.output_type, block_state.mask_overlay_kwargs)
+
+        if block_state.mask_overlay_kwargs is None:
+            mask_overlay_kwargs = {}
+        else:
+            mask_overlay_kwargs = block_state.mask_overlay_kwargs
+
+        block_state.images = components.image_mask_processor.postprocess(
+            image=block_state.images,
+            **mask_overlay_kwargs,
+        )
+
+        self.set_block_state(state, block_state)
+        return components, state

denoise.py

@@ -0,0 +1,407 @@
+# 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.
+
+import inspect
+
+import torch
+
+from diffusers.configuration_utils import FrozenDict
+from diffusers.guiders import ClassifierFreeGuidance
+from .transformer_krea2 import Krea2Transformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import logging
+from diffusers.modular_pipelines.modular_pipeline import BlockState, LoopSequentialPipelineBlocks, ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+from .modular_pipeline import Krea2ModularPipeline
+
+
+logger = logging.get_logger(__name__)
+
+# ====================
+# 1. LOOP STEPS (run at each denoising step)
+# ====================
+
+
+# loop step:before denoiser
+class Krea2LoopBeforeDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that prepares the latent input for the denoiser. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `Krea2DenoiseLoopWrapper`)"
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        # one timestep
+        block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype)
+        block_state.latent_model_input = block_state.latents
+        return components, block_state
+
+
+# loop step:before denoiser (edit) -- appends the clean reference tokens to the denoiser input each step
+class Krea2EditLoopBeforeDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that prepares the latent input for the edit denoiser: it appends the "
+            "packed clean reference tokens after the noisy image tokens. This block should be used to compose the "
+            "`sub_blocks` attribute of a `LoopSequentialPipelineBlocks` object (e.g. `Krea2EditDenoiseStep`)."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial latents to use for the denoising process. Can be generated in prepare_latent step.",
+            ),
+            InputParam(
+                name="reference_latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="Packed clean reference tokens to append to the denoiser sequence. Can be generated in the reference latents step.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        block_state.timestep = t.expand(block_state.latents.shape[0]).to(block_state.latents.dtype)
+        # Reference tokens are shared across the batch; expand and append them after the noisy image tokens.
+        reference_latents = block_state.reference_latents.expand(block_state.latents.shape[0], -1, -1)
+        block_state.latent_model_input = torch.cat(
+            [block_state.latents, reference_latents.to(block_state.latents.dtype)], dim=1
+        )
+        return components, block_state
+
+
+# loop step:denoiser
+class Krea2LoopDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that denoise the latent input for the denoiser. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `Krea2DenoiseLoopWrapper`)"
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "guider",
+                ClassifierFreeGuidance,
+                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.template("denoiser_input_fields"),
+            InputParam(
+                "position_ids",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The rotary coordinates for the combined text-image sequence. Can be generated in prepare_rope_inputs step.",
+            ),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        guider_inputs = {
+            "encoder_hidden_states": (
+                getattr(block_state, "prompt_embeds", None),
+                getattr(block_state, "negative_prompt_embeds", None),
+            ),
+            "encoder_attention_mask": (
+                getattr(block_state, "prompt_embeds_mask", None),
+                getattr(block_state, "negative_prompt_embeds_mask", None),
+            ),
+        }
+
+        transformer_args = set(inspect.signature(components.transformer.forward).parameters.keys())
+        additional_cond_kwargs = {}
+        for field_name, field_value in block_state.denoiser_input_fields.items():
+            if field_name in transformer_args and field_name not in guider_inputs:
+                additional_cond_kwargs[field_name] = field_value
+        block_state.additional_cond_kwargs.update(additional_cond_kwargs)
+
+        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(components.transformer)
+            cond_kwargs = {input_name: getattr(guider_state_batch, input_name) for input_name in guider_inputs.keys()}
+
+            guider_state_batch.noise_pred = components.transformer(
+                hidden_states=block_state.latent_model_input,
+                timestep=block_state.timestep / 1000,
+                return_dict=False,
+                **cond_kwargs,
+                **block_state.additional_cond_kwargs,
+            )[0]
+
+            components.guider.cleanup_models(components.transformer)
+
+        guider_output = components.guider(guider_state)
+        block_state.noise_pred = guider_output.pred
+
+        return components, block_state
+
+
+# loop step:after denoiser
+class Krea2LoopAfterDenoiser(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that updates the latents. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `Krea2DenoiseLoopWrapper`)"
+        )
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam.template("latents"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        latents_dtype = block_state.latents.dtype
+        block_state.latents = components.scheduler.step(
+            block_state.noise_pred,
+            t,
+            block_state.latents,
+            return_dict=False,
+        )[0]
+
+        if block_state.latents.dtype != latents_dtype:
+            if torch.backends.mps.is_available():
+                # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                block_state.latents = block_state.latents.to(latents_dtype)
+
+        return components, block_state
+
+
+class Krea2LoopAfterDenoiserInpaint(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "step within the denoising loop that updates the latents using mask and image_latents for inpainting. "
+            "This block should be used to compose the `sub_blocks` attribute of a `LoopSequentialPipelineBlocks` "
+            "object (e.g. `Krea2DenoiseLoopWrapper`)"
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                "mask",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The mask to use for the inpainting process. Can be generated in inpaint prepare latents step.",
+            ),
+            InputParam.template("image_latents"),
+            InputParam(
+                "initial_noise",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The initial noise to use for the inpainting process. Can be generated in inpaint prepare latents step.",
+            ),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam.template("latents"),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, block_state: BlockState, i: int, t: torch.Tensor):
+        block_state.init_latents_proper = block_state.image_latents
+        if i < len(block_state.timesteps) - 1:
+            block_state.noise_timestep = block_state.timesteps[i + 1]
+            block_state.init_latents_proper = components.scheduler.scale_noise(
+                block_state.init_latents_proper, torch.tensor([block_state.noise_timestep]), block_state.initial_noise
+            )
+
+        block_state.latents = (
+            1 - block_state.mask
+        ) * block_state.init_latents_proper + block_state.mask * block_state.latents
+
+        return components, block_state
+
+
+# ====================
+# 2. DENOISE LOOP WRAPPER: define the denoising loop logic
+# ====================
+class Krea2DenoiseLoopWrapper(LoopSequentialPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Pipeline block that iteratively denoise the latents over `timesteps`. "
+            "The specific steps with each iteration can be customized with `sub_blocks` attributes"
+        )
+
+    @property
+    def loop_expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler),
+        ]
+
+    @property
+    def loop_inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="timesteps",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The timesteps to use for the denoising process. Can be generated in set_timesteps step.",
+            ),
+            InputParam.template("num_inference_steps", required=True),
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        block_state.num_warmup_steps = max(
+            len(block_state.timesteps) - block_state.num_inference_steps * components.scheduler.order, 0
+        )
+
+        block_state.additional_cond_kwargs = {}
+
+        with self.progress_bar(total=block_state.num_inference_steps) as progress_bar:
+            for i, t in enumerate(block_state.timesteps):
+                components, block_state = self.loop_step(components, block_state, i=i, t=t)
+                if i == len(block_state.timesteps) - 1 or (
+                    (i + 1) > block_state.num_warmup_steps and (i + 1) % components.scheduler.order == 0
+                ):
+                    progress_bar.update()
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+# ====================
+# 3. DENOISE STEPS: compose the denoising loop with loop wrapper + loop steps
+# ====================
+
+
+# Krea 2 (text2image, image2image)
+class Krea2DenoiseStep(Krea2DenoiseLoopWrapper):
+    model_name = "krea2"
+
+    block_classes = [
+        Krea2LoopBeforeDenoiser,
+        Krea2LoopDenoiser,
+        Krea2LoopAfterDenoiser,
+    ]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoise the latents.\n"
+            "Its loop logic is defined in `Krea2DenoiseLoopWrapper.__call__` method\n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+            " - `Krea2LoopBeforeDenoiser`\n"
+            " - `Krea2LoopDenoiser`\n"
+            " - `Krea2LoopAfterDenoiser`\n"
+            "This block supports text2image and image2image tasks for Krea 2."
+        )
+
+
+# Krea 2 (inpainting)
+class Krea2InpaintDenoiseStep(Krea2DenoiseLoopWrapper):
+    model_name = "krea2"
+    block_classes = [
+        Krea2LoopBeforeDenoiser,
+        Krea2LoopDenoiser,
+        Krea2LoopAfterDenoiser,
+        Krea2LoopAfterDenoiserInpaint,
+    ]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser", "after_denoiser_inpaint"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoise the latents. \n"
+            "Its loop logic is defined in `Krea2DenoiseLoopWrapper.__call__` method \n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+            " - `Krea2LoopBeforeDenoiser`\n"
+            " - `Krea2LoopDenoiser`\n"
+            " - `Krea2LoopAfterDenoiser`\n"
+            " - `Krea2LoopAfterDenoiserInpaint`\n"
+            "This block supports inpainting tasks for Krea 2."
+        )
+
+
+# Krea 2 (reference-image edit)
+class Krea2EditDenoiseStep(Krea2DenoiseLoopWrapper):
+    model_name = "krea2"
+    block_classes = [
+        Krea2EditLoopBeforeDenoiser,
+        Krea2LoopDenoiser,
+        Krea2LoopAfterDenoiser,
+    ]
+    block_names = ["before_denoiser", "denoiser", "after_denoiser"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Denoise step that iteratively denoise the latents for the reference-image edit task.\n"
+            "Its loop logic is defined in `Krea2DenoiseLoopWrapper.__call__` method\n"
+            "At each iteration, it runs blocks defined in `sub_blocks` sequencially:\n"
+            " - `Krea2EditLoopBeforeDenoiser` (appends the clean reference tokens)\n"
+            " - `Krea2LoopDenoiser`\n"
+            " - `Krea2LoopAfterDenoiser`\n"
+            "This block supports reference-image (edit) generation for Krea 2."
+        )

encoders.py

@@ -0,0 +1,844 @@
+# 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

inputs.py

@@ -0,0 +1,406 @@
+# 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.
+
+
+import torch
+
+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
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.inputs.repeat_tensor_to_batch_size
+def repeat_tensor_to_batch_size(
+    input_name: str,
+    input_tensor: torch.Tensor,
+    batch_size: int,
+    num_images_per_prompt: int = 1,
+) -> torch.Tensor:
+    """Repeat tensor elements to match the final batch size.
+
+    This function expands a tensor's batch dimension to match the final batch size (batch_size * num_images_per_prompt)
+    by repeating each element along dimension 0.
+
+    The input tensor must have batch size 1 or batch_size. The function will:
+    - If batch size is 1: repeat each element (batch_size * num_images_per_prompt) times
+    - If batch size equals batch_size: repeat each element num_images_per_prompt times
+
+    Args:
+        input_name (str): Name of the input tensor (used for error messages)
+        input_tensor (torch.Tensor): The tensor to repeat. Must have batch size 1 or batch_size.
+        batch_size (int): The base batch size (number of prompts)
+        num_images_per_prompt (int, optional): Number of images to generate per prompt. Defaults to 1.
+
+    Returns:
+        torch.Tensor: The repeated tensor with final batch size (batch_size * num_images_per_prompt)
+
+    Raises:
+        ValueError: If input_tensor is not a torch.Tensor or has invalid batch size
+
+    Examples:
+        tensor = torch.tensor([[1, 2, 3]]) # shape: [1, 3] repeated = repeat_tensor_to_batch_size("image", tensor,
+        batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3], [1, 2, 3]]) - shape:
+        [4, 3]
+
+        tensor = torch.tensor([[1, 2, 3], [4, 5, 6]]) # shape: [2, 3] repeated = repeat_tensor_to_batch_size("image",
+        tensor, batch_size=2, num_images_per_prompt=2) repeated # tensor([[1, 2, 3], [1, 2, 3], [4, 5, 6], [4, 5, 6]])
+        - shape: [4, 3]
+    """
+    # make sure input is a tensor
+    if not isinstance(input_tensor, torch.Tensor):
+        raise ValueError(f"`{input_name}` must be a tensor")
+
+    # make sure input tensor e.g. image_latents has batch size 1 or batch_size same as prompts
+    if input_tensor.shape[0] == 1:
+        repeat_by = batch_size * num_images_per_prompt
+    elif input_tensor.shape[0] == batch_size:
+        repeat_by = num_images_per_prompt
+    else:
+        raise ValueError(
+            f"`{input_name}` must have have batch size 1 or {batch_size}, but got {input_tensor.shape[0]}"
+        )
+
+    # expand the tensor to match the batch_size * num_images_per_prompt
+    input_tensor = input_tensor.repeat_interleave(repeat_by, dim=0)
+
+    return input_tensor
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.inputs.calculate_dimension_from_latents
+def calculate_dimension_from_latents(latents: torch.Tensor, vae_scale_factor: int) -> tuple[int, int]:
+    """Calculate image dimensions from latent tensor dimensions.
+
+    This function converts latent space dimensions to image space dimensions by multiplying the latent height and width
+    by the VAE scale factor.
+
+    Args:
+        latents (torch.Tensor): The latent tensor. Must have 4 or 5 dimensions.
+            Expected shapes: [batch, channels, height, width] or [batch, channels, frames, height, width]
+        vae_scale_factor (int): The scale factor used by the VAE to compress images.
+            Typically 8 for most VAEs (image is 8x larger than latents in each dimension)
+
+    Returns:
+        tuple[int, int]: The calculated image dimensions as (height, width)
+
+    Raises:
+        ValueError: If latents tensor doesn't have 4 or 5 dimensions
+
+    """
+    # make sure the latents are not packed
+    if latents.ndim != 4 and latents.ndim != 5:
+        raise ValueError(f"unpacked latents must have 4 or 5 dimensions, but got {latents.ndim}")
+
+    latent_height, latent_width = latents.shape[-2:]
+
+    height = latent_height * vae_scale_factor
+    width = latent_width * vae_scale_factor
+
+    return height, width
+
+
+class Krea2TextInputsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        summary_section = (
+            "Text input processing step that standardizes text embeddings for the pipeline.\n"
+            "This step:\n"
+            "  1. Determines `batch_size` and `dtype` based on `prompt_embeds`\n"
+            "  2. Ensures all text embeddings have consistent batch sizes (batch_size * num_images_per_prompt)"
+        )
+
+        # Placement guidance
+        placement_section = "\n\nThis block should be placed after all encoder steps to process the text embeddings before they are used in subsequent pipeline steps."
+
+        return summary_section + placement_section
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam.template("num_images_per_prompt"),
+            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="The batch size of the prompt embeddings"),
+            OutputParam(name="dtype", type_hint=torch.dtype, description="The data type of the prompt embeddings"),
+            OutputParam.template("prompt_embeds", note="batch-expanded"),
+            OutputParam.template("prompt_embeds_mask", note="batch-expanded"),
+            OutputParam.template("negative_prompt_embeds", note="batch-expanded"),
+            OutputParam.template("negative_prompt_embeds_mask", note="batch-expanded"),
+        ]
+
+    @staticmethod
+    def check_inputs(
+        prompt_embeds,
+        prompt_embeds_mask,
+        negative_prompt_embeds,
+        negative_prompt_embeds_mask,
+    ):
+        if prompt_embeds.ndim != 4:
+            raise ValueError(
+                f"`prompt_embeds` must have 4 dimensions (batch_size, text_seq_len, num_text_layers, text_hidden_dim), but got {prompt_embeds.ndim}"
+            )
+
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            raise ValueError("`negative_prompt_embeds_mask` is required when `negative_prompt_embeds` is not None")
+
+        if negative_prompt_embeds is None and negative_prompt_embeds_mask is not None:
+            raise ValueError("cannot pass `negative_prompt_embeds_mask` without `negative_prompt_embeds`")
+
+        if prompt_embeds_mask.shape[0] != prompt_embeds.shape[0]:
+            raise ValueError("`prompt_embeds_mask` must have the same batch size as `prompt_embeds`")
+
+        elif negative_prompt_embeds is not None and negative_prompt_embeds.shape[0] != prompt_embeds.shape[0]:
+            raise ValueError("`negative_prompt_embeds` must have the same batch size as `prompt_embeds`")
+
+        elif (
+            negative_prompt_embeds_mask is not None and negative_prompt_embeds_mask.shape[0] != prompt_embeds.shape[0]
+        ):
+            raise ValueError("`negative_prompt_embeds_mask` must have the same batch size as `prompt_embeds`")
+
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        self.check_inputs(
+            prompt_embeds=block_state.prompt_embeds,
+            prompt_embeds_mask=block_state.prompt_embeds_mask,
+            negative_prompt_embeds=block_state.negative_prompt_embeds,
+            negative_prompt_embeds_mask=block_state.negative_prompt_embeds_mask,
+        )
+
+        block_state.batch_size = block_state.prompt_embeds.shape[0]
+        block_state.dtype = block_state.prompt_embeds.dtype
+
+        # Krea 2 prompt embeddings are 4D: (batch_size, text_seq_len, num_text_layers, text_hidden_dim)
+        _, seq_len, num_text_layers, dim = block_state.prompt_embeds.shape
+
+        block_state.prompt_embeds = block_state.prompt_embeds.repeat(1, block_state.num_images_per_prompt, 1, 1)
+        block_state.prompt_embeds = block_state.prompt_embeds.view(
+            block_state.batch_size * block_state.num_images_per_prompt, seq_len, num_text_layers, dim
+        )
+
+        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.repeat(1, block_state.num_images_per_prompt)
+        block_state.prompt_embeds_mask = block_state.prompt_embeds_mask.view(
+            block_state.batch_size * block_state.num_images_per_prompt, seq_len
+        )
+
+        if block_state.negative_prompt_embeds is not None:
+            _, seq_len, num_text_layers, dim = block_state.negative_prompt_embeds.shape
+            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.repeat(
+                1, block_state.num_images_per_prompt, 1, 1
+            )
+            block_state.negative_prompt_embeds = block_state.negative_prompt_embeds.view(
+                block_state.batch_size * block_state.num_images_per_prompt, seq_len, num_text_layers, dim
+            )
+
+            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.repeat(
+                1, block_state.num_images_per_prompt
+            )
+            block_state.negative_prompt_embeds_mask = block_state.negative_prompt_embeds_mask.view(
+                block_state.batch_size * block_state.num_images_per_prompt, seq_len
+            )
+
+        self.set_block_state(state, block_state)
+
+        return components, state
+
+
+class Krea2AdditionalInputsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    def __init__(
+        self,
+        image_latent_inputs: list[InputParam] | None = None,
+        additional_batch_inputs: list[InputParam] | None = None,
+    ):
+        """Input processing step for additional (non-text) inputs.
+
+        For image latent inputs: updates height/width if None, patchifies, and expands batch size. For additional batch
+        inputs: expands batch dimensions to match the final batch size.
+
+        Args:
+            image_latent_inputs (list[InputParam], optional): Image latent inputs to process. Defaults to
+                `[InputParam.template("image_latents")]`.
+            additional_batch_inputs (list[InputParam], optional): Additional tensor inputs that only need batch
+                expansion. Defaults to `[]`.
+        """
+        # by default, process `image_latents`
+        if image_latent_inputs is None:
+            image_latent_inputs = [InputParam.template("image_latents")]
+        if additional_batch_inputs is None:
+            additional_batch_inputs = []
+
+        if not isinstance(image_latent_inputs, list):
+            raise ValueError(f"image_latent_inputs must be a list, but got {type(image_latent_inputs)}")
+        else:
+            for input_param in image_latent_inputs:
+                if not isinstance(input_param, InputParam):
+                    raise ValueError(f"image_latent_inputs must be a list of InputParam, but got {type(input_param)}")
+
+        if not isinstance(additional_batch_inputs, list):
+            raise ValueError(f"additional_batch_inputs must be a list, but got {type(additional_batch_inputs)}")
+        else:
+            for input_param in additional_batch_inputs:
+                if not isinstance(input_param, InputParam):
+                    raise ValueError(
+                        f"additional_batch_inputs must be a list of InputParam, but got {type(input_param)}"
+                    )
+
+        self._image_latent_inputs = image_latent_inputs
+        self._additional_batch_inputs = additional_batch_inputs
+        super().__init__()
+
+    @property
+    def description(self) -> str:
+        summary_section = (
+            "Input processing step that:\n"
+            "  1. For image latent inputs: Updates height/width if None, patchifies, and expands batch size\n"
+            "  2. For additional batch inputs: Expands batch dimensions to match final batch size"
+        )
+
+        inputs_info = ""
+        if self._image_latent_inputs or self._additional_batch_inputs:
+            inputs_info = "\n\nConfigured inputs:"
+            if self._image_latent_inputs:
+                inputs_info += f"\n  - Image latent inputs: {[p.name for p in self._image_latent_inputs]}"
+            if self._additional_batch_inputs:
+                inputs_info += f"\n  - Additional batch inputs: {[p.name for p in self._additional_batch_inputs]}"
+
+        placement_section = "\n\nThis block should be placed after the encoder steps and the text input step."
+
+        return summary_section + inputs_info + placement_section
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec("pachifier", Krea2Pachifier, default_creation_method="from_config"),
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        inputs = [
+            InputParam.template("num_images_per_prompt"),
+            InputParam.template("batch_size"),
+            InputParam.template("height"),
+            InputParam.template("width"),
+        ]
+        # default is `image_latents`
+        inputs += self._image_latent_inputs + self._additional_batch_inputs
+
+        return inputs
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        outputs = [
+            OutputParam(
+                name="image_height",
+                type_hint=int,
+                description="The image height calculated from the image latents dimension",
+            ),
+            OutputParam(
+                name="image_width",
+                type_hint=int,
+                description="The image width calculated from the image latents dimension",
+            ),
+        ]
+
+        # `height`/`width` are not new outputs, but they will be updated if any image latent inputs are provided
+        if len(self._image_latent_inputs) > 0:
+            outputs.append(
+                OutputParam(name="height", type_hint=int, description="if not provided, updated to image height")
+            )
+            outputs.append(
+                OutputParam(name="width", type_hint=int, description="if not provided, updated to image width")
+            )
+
+        # image latent inputs are modified in place (patchified and batch-expanded)
+        for input_param in self._image_latent_inputs:
+            outputs.append(
+                OutputParam(
+                    name=input_param.name,
+                    type_hint=input_param.type_hint,
+                    description=input_param.description + " (patchified and batch-expanded)",
+                )
+            )
+
+        # additional batch inputs (batch-expanded only)
+        for input_param in self._additional_batch_inputs:
+            outputs.append(
+                OutputParam(
+                    name=input_param.name,
+                    type_hint=input_param.type_hint,
+                    description=input_param.description + " (batch-expanded)",
+                )
+            )
+
+        return outputs
+
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+
+        # Process image latent inputs
+        for input_param in self._image_latent_inputs:
+            image_latent_input_name = input_param.name
+            image_latent_tensor = getattr(block_state, image_latent_input_name)
+            if image_latent_tensor is None:
+                continue
+
+            # 1. Calculate height/width from latents and update if not provided
+            height, width = calculate_dimension_from_latents(image_latent_tensor, components.vae_scale_factor)
+            block_state.height = block_state.height or height
+            block_state.width = block_state.width or width
+
+            if not hasattr(block_state, "image_height"):
+                block_state.image_height = height
+            if not hasattr(block_state, "image_width"):
+                block_state.image_width = width
+
+            # 2. Patchify
+            image_latent_tensor = components.pachifier.pack_latents(image_latent_tensor)
+
+            # 3. Expand batch size
+            image_latent_tensor = repeat_tensor_to_batch_size(
+                input_name=image_latent_input_name,
+                input_tensor=image_latent_tensor,
+                num_images_per_prompt=block_state.num_images_per_prompt,
+                batch_size=block_state.batch_size,
+            )
+
+            setattr(block_state, image_latent_input_name, image_latent_tensor)
+
+        # Process additional batch inputs (only batch expansion)
+        for input_param in self._additional_batch_inputs:
+            input_name = input_param.name
+            input_tensor = getattr(block_state, input_name)
+            if input_tensor is None:
+                continue
+
+            input_tensor = repeat_tensor_to_batch_size(
+                input_name=input_name,
+                input_tensor=input_tensor,
+                num_images_per_prompt=block_state.num_images_per_prompt,
+                batch_size=block_state.batch_size,
+            )
+
+            setattr(block_state, input_name, input_tensor)
+
+        self.set_block_state(state, block_state)
+        return components, state

modular_blocks_krea2.py

@@ -0,0 +1,879 @@
+# 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.
+
+import torch
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines.modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import InputParam, InsertableDict, OutputParam
+from .before_denoise import (
+    Krea2CreateMaskLatentsStep,
+    Krea2PrepareLatentsStep,
+    Krea2PrepareLatentsWithStrengthStep,
+    Krea2RoPEInputsStep,
+    Krea2SetTimestepsStep,
+    Krea2SetTimestepsWithStrengthStep,
+)
+from .decoders import (
+    Krea2AfterDenoiseStep,
+    Krea2DecoderStep,
+    Krea2InpaintProcessImagesOutputStep,
+    Krea2ProcessImagesOutputStep,
+)
+from .denoise import (
+    Krea2DenoiseStep,
+    Krea2InpaintDenoiseStep,
+)
+from .encoders import (
+    Krea2InpaintProcessImagesInputStep,
+    Krea2ProcessImagesInputStep,
+    Krea2TextEncoderStep,
+    Krea2VaeEncoderStep,
+)
+from .inputs import (
+    Krea2AdditionalInputsStep,
+    Krea2TextInputsStep,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+# ====================
+# 1. TEXT ENCODER
+# ====================
+
+
+# auto_docstring
+class Krea2AutoTextEncoderStep(AutoPipelineBlocks):
+    """
+    Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block.
+       - `Krea2TextEncoderStep` (text_encoder) is used when `prompt` is provided.
+       - if `prompt` is not provided, step will be skipped.
+
+      Components:
+          text_encoder (`Qwen3VLModel`): The text encoder to use tokenizer (`Qwen2Tokenizer`): The tokenizer to use
+          guider (`ClassifierFreeGuidance`)
+
+      Inputs:
+          prompt (`str`, *optional*):
+              The prompt or prompts to guide image generation.
+          negative_prompt (`str`, *optional*):
+              The prompt or prompts not to guide the image generation.
+          max_sequence_length (`int`, *optional*, defaults to 512):
+              Maximum sequence length for prompt encoding.
+
+      Outputs:
+          prompt_embeds (`Tensor`):
+              The prompt embeddings.
+          prompt_embeds_mask (`Tensor`):
+              The encoder attention mask.
+          negative_prompt_embeds (`Tensor`):
+              The negative prompt embeddings.
+          negative_prompt_embeds_mask (`Tensor`):
+              The negative prompt embeddings mask.
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2TextEncoderStep()]
+    block_names = ["text_encoder"]
+    block_trigger_inputs = ["prompt"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "Text encoder step that encodes the text prompt into a text embedding. This is an auto pipeline block.\n"
+            " - `Krea2TextEncoderStep` (text_encoder) is used when `prompt` is provided.\n"
+            " - if `prompt` is not provided, step will be skipped."
+        )
+
+
+# ====================
+# 2. VAE ENCODER
+# ====================
+
+
+# auto_docstring
+class Krea2InpaintVaeEncoderStep(SequentialPipelineBlocks):
+    """
+    This step is used for processing image and mask inputs for inpainting tasks. It:
+       - Resizes the image to the target size, based on `height` and `width`.
+       - Processes and updates `image` and `mask_image`.
+       - Creates `image_latents`.
+
+      Components:
+          image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLQwenImage`)
+
+      Inputs:
+          mask_image (`Image`):
+              Mask image for inpainting.
+          image (`Image | list`):
+              Reference image(s) for denoising. Can be a single image or list of images.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          padding_mask_crop (`int`, *optional*):
+              Padding for mask cropping in inpainting.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+
+      Outputs:
+          processed_image (`Tensor`):
+              The processed image
+          processed_mask_image (`Tensor`):
+              The processed mask image
+          mask_overlay_kwargs (`dict`):
+              The kwargs for the postprocess step to apply the mask overlay
+          image_latents (`Tensor`):
+              The latent representation of the input image.
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2InpaintProcessImagesInputStep(), Krea2VaeEncoderStep()]
+    block_names = ["preprocess", "encode"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "This step is used for processing image and mask inputs for inpainting tasks. It:\n"
+            " - Resizes the image to the target size, based on `height` and `width`.\n"
+            " - Processes and updates `image` and `mask_image`.\n"
+            " - Creates `image_latents`."
+        )
+
+
+# auto_docstring
+class Krea2Img2ImgVaeEncoderStep(SequentialPipelineBlocks):
+    """
+    Vae encoder step that preprocess and encode the image inputs into their latent representations.
+
+      Components:
+          image_processor (`VaeImageProcessor`) vae (`AutoencoderKLQwenImage`)
+
+      Inputs:
+          image (`Image | list`):
+              Reference image(s) for denoising. Can be a single image or list of images.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+
+      Outputs:
+          processed_image (`Tensor`):
+              The processed image
+          image_latents (`Tensor`):
+              The latent representation of the input image.
+    """
+
+    model_name = "krea2"
+
+    block_classes = [Krea2ProcessImagesInputStep(), Krea2VaeEncoderStep()]
+    block_names = ["preprocess", "encode"]
+
+    @property
+    def description(self) -> str:
+        return "Vae encoder step that preprocess and encode the image inputs into their latent representations."
+
+
+class Krea2AutoVaeEncoderStep(AutoPipelineBlocks):
+    model_name = "krea2"
+    block_classes = [Krea2InpaintVaeEncoderStep, Krea2Img2ImgVaeEncoderStep]
+    block_names = ["inpaint", "img2img"]
+    block_trigger_inputs = ["mask_image", "image"]
+
+    @property
+    def description(self):
+        return (
+            "Vae encoder step that encode the image inputs into their latent representations.\n"
+            + "This is an auto pipeline block.\n"
+            + " - `Krea2InpaintVaeEncoderStep` (inpaint) is used when `mask_image` is provided.\n"
+            + " - `Krea2Img2ImgVaeEncoderStep` (img2img) is used when `image` is provided.\n"
+            + " - if `mask_image` or `image` is not provided, step will be skipped."
+        )
+
+
+# ====================
+# 3. DENOISE (input -> prepare_latents -> set_timesteps -> prepare_rope_inputs -> denoise -> after_denoise)
+# ====================
+
+
+# assemble input steps
+# auto_docstring
+class Krea2Img2ImgInputStep(SequentialPipelineBlocks):
+    """
+    Input step that prepares the inputs for the img2img denoising step. It:
+       - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).
+       - update height/width based `image_latents`, patchify `image_latents`.
+
+      Components:
+          pachifier (`Krea2Pachifier`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          image_latents (`Tensor`):
+              image latents used to guide the image generation. Can be generated from vae_encoder step.
+
+      Outputs:
+          batch_size (`int`):
+              The batch size of the prompt embeddings
+          dtype (`dtype`):
+              The data type of the prompt embeddings
+          prompt_embeds (`Tensor`):
+              The prompt embeddings. (batch-expanded)
+          prompt_embeds_mask (`Tensor`):
+              The encoder attention mask. (batch-expanded)
+          negative_prompt_embeds (`Tensor`):
+              The negative prompt embeddings. (batch-expanded)
+          negative_prompt_embeds_mask (`Tensor`):
+              The negative prompt embeddings mask. (batch-expanded)
+          image_height (`int`):
+              The image height calculated from the image latents dimension
+          image_width (`int`):
+              The image width calculated from the image latents dimension
+          height (`int`):
+              if not provided, updated to image height
+          width (`int`):
+              if not provided, updated to image width
+          image_latents (`Tensor`):
+              image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+              batch-expanded)
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2TextInputsStep(), Krea2AdditionalInputsStep()]
+    block_names = ["text_inputs", "additional_inputs"]
+
+    @property
+    def description(self):
+        return (
+            "Input step that prepares the inputs for the img2img denoising step. It:\n"
+            " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents`).\n"
+            " - update height/width based `image_latents`, patchify `image_latents`."
+        )
+
+
+# auto_docstring
+class Krea2InpaintInputStep(SequentialPipelineBlocks):
+    """
+    Input step that prepares the inputs for the inpainting denoising step. It:
+       - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents` and
+         `processed_mask_image`).
+       - update height/width based `image_latents`, patchify `image_latents`.
+
+      Components:
+          pachifier (`Krea2Pachifier`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          image_latents (`Tensor`, *optional*):
+              image latents used to guide the image generation. Can be generated from vae_encoder step.
+          processed_mask_image (`Tensor`, *optional*):
+              The processed mask image
+
+      Outputs:
+          batch_size (`int`):
+              The batch size of the prompt embeddings
+          dtype (`dtype`):
+              The data type of the prompt embeddings
+          prompt_embeds (`Tensor`):
+              The prompt embeddings. (batch-expanded)
+          prompt_embeds_mask (`Tensor`):
+              The encoder attention mask. (batch-expanded)
+          negative_prompt_embeds (`Tensor`):
+              The negative prompt embeddings. (batch-expanded)
+          negative_prompt_embeds_mask (`Tensor`):
+              The negative prompt embeddings mask. (batch-expanded)
+          image_height (`int`):
+              The image height calculated from the image latents dimension
+          image_width (`int`):
+              The image width calculated from the image latents dimension
+          height (`int`):
+              if not provided, updated to image height
+          width (`int`):
+              if not provided, updated to image width
+          image_latents (`Tensor`):
+              image latents used to guide the image generation. Can be generated from vae_encoder step. (patchified and
+              batch-expanded)
+          processed_mask_image (`Tensor`):
+              The processed mask image (batch-expanded)
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2TextInputsStep(),
+        Krea2AdditionalInputsStep(
+            additional_batch_inputs=[
+                InputParam(name="processed_mask_image", type_hint=torch.Tensor, description="The processed mask image")
+            ]
+        ),
+    ]
+    block_names = ["text_inputs", "additional_inputs"]
+
+    @property
+    def description(self):
+        return (
+            "Input step that prepares the inputs for the inpainting denoising step. It:\n"
+            " - make sure the text embeddings have consistent batch size as well as the additional inputs (`image_latents` and `processed_mask_image`).\n"
+            " - update height/width based `image_latents`, patchify `image_latents`."
+        )
+
+
+# assemble prepare latents steps
+# auto_docstring
+class Krea2InpaintPrepareLatentsStep(SequentialPipelineBlocks):
+    """
+    This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:
+       - Add noise to the image latents to create the latents input for the denoiser.
+       - Create the patchified latents `mask` based on the processed mask image.
+
+      Components:
+          scheduler (`FlowMatchEulerDiscreteScheduler`) pachifier (`Krea2Pachifier`)
+
+      Inputs:
+          latents (`Tensor`):
+              The initial random noised, can be generated in prepare latent step.
+          image_latents (`Tensor`):
+              image latents used to guide the image generation. Can be generated from vae_encoder step. (Can be
+              generated from vae encoder and updated in input step.)
+          timesteps (`Tensor`):
+              The timesteps to use for the denoising process. Can be generated in set_timesteps step.
+          processed_mask_image (`Tensor`):
+              The processed mask to use for the inpainting process.
+          height (`int`):
+              The height in pixels of the generated image.
+          width (`int`):
+              The width in pixels of the generated image.
+          dtype (`dtype`, *optional*, defaults to torch.float32):
+              The dtype of the model inputs, can be generated in input step.
+
+      Outputs:
+          initial_noise (`Tensor`):
+              The initial random noised used for inpainting denoising.
+          latents (`Tensor`):
+              The scaled noisy latents to use for inpainting/image-to-image denoising.
+          mask (`Tensor`):
+              The mask to use for the inpainting process.
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2PrepareLatentsWithStrengthStep(), Krea2CreateMaskLatentsStep()]
+    block_names = ["add_noise_to_latents", "create_mask_latents"]
+
+    @property
+    def description(self) -> str:
+        return (
+            "This step prepares the latents/image_latents and mask inputs for the inpainting denoising step. It:\n"
+            " - Add noise to the image latents to create the latents input for the denoiser.\n"
+            " - Create the patchified latents `mask` based on the processed mask image.\n"
+        )
+
+
+# assemble denoising steps
+
+
+# Krea 2 (text2image)
+# auto_docstring
+class Krea2CoreDenoiseStep(SequentialPipelineBlocks):
+    """
+    step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs
+    (timesteps, latents, rope inputs etc.).
+
+      Components:
+          pachifier (`Krea2Pachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+          transformer (`Krea2Transformer2DModel`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          latents (`Tensor`, *optional*):
+              Pre-generated noisy latents for image generation.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`, *optional*, defaults to 28):
+              The number of denoising steps.
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process.
+          mu (`float`, *optional*):
+              Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if
+              not provided, computed from the image sequence length (base checkpoint behavior).
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2TextInputsStep(),
+        Krea2PrepareLatentsStep(),
+        Krea2SetTimestepsStep(),
+        Krea2RoPEInputsStep(),
+        Krea2DenoiseStep(),
+        Krea2AfterDenoiseStep(),
+    ]
+    block_names = [
+        "input",
+        "prepare_latents",
+        "set_timesteps",
+        "prepare_rope_inputs",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self):
+        return "step that denoise noise into image for text2image task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# Krea 2 (inpainting)
+# auto_docstring
+class Krea2InpaintCoreDenoiseStep(SequentialPipelineBlocks):
+    """
+    step that denoise noise into image for inpaint task. It includes the denoise loop, as well as prepare the inputs
+    (timesteps, latents, rope inputs etc.).
+
+      Components:
+          pachifier (`Krea2Pachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+          transformer (`Krea2Transformer2DModel`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          image_latents (`Tensor`, *optional*):
+              image latents used to guide the image generation. Can be generated from vae_encoder step.
+          processed_mask_image (`Tensor`, *optional*):
+              The processed mask image
+          latents (`Tensor`, *optional*):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`, *optional*, defaults to 28):
+              The number of denoising steps.
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process.
+          mu (`float`, *optional*):
+              Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if
+              not provided, computed from the image sequence length (base checkpoint behavior).
+          strength (`float`, *optional*, defaults to 0.9):
+              Strength for img2img/inpainting.
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2InpaintInputStep(),
+        Krea2PrepareLatentsStep(),
+        Krea2SetTimestepsWithStrengthStep(),
+        Krea2InpaintPrepareLatentsStep(),
+        Krea2RoPEInputsStep(),
+        Krea2InpaintDenoiseStep(),
+        Krea2AfterDenoiseStep(),
+    ]
+    block_names = [
+        "input",
+        "prepare_latents",
+        "set_timesteps",
+        "prepare_inpaint_latents",
+        "prepare_rope_inputs",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self):
+        return "step that denoise noise into image for inpaint task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# Krea 2 (image2image)
+# auto_docstring
+class Krea2Img2ImgCoreDenoiseStep(SequentialPipelineBlocks):
+    """
+    step that denoise noise into image for img2img task. It includes the denoise loop, as well as prepare the inputs
+    (timesteps, latents, rope inputs etc.).
+
+      Components:
+          pachifier (`Krea2Pachifier`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`ClassifierFreeGuidance`)
+          transformer (`Krea2Transformer2DModel`)
+
+      Inputs:
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          image_latents (`Tensor`):
+              image latents used to guide the image generation. Can be generated from vae_encoder step.
+          latents (`Tensor`, *optional*):
+              Pre-generated noisy latents for image generation.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_inference_steps (`int`, *optional*, defaults to 28):
+              The number of denoising steps.
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process.
+          mu (`float`, *optional*):
+              Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if
+              not provided, computed from the image sequence length (base checkpoint behavior).
+          strength (`float`, *optional*, defaults to 0.9):
+              Strength for img2img/inpainting.
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+
+      Outputs:
+          latents (`Tensor`):
+              Denoised latents.
+    """
+
+    model_name = "krea2"
+    block_classes = [
+        Krea2Img2ImgInputStep(),
+        Krea2PrepareLatentsStep(),
+        Krea2SetTimestepsWithStrengthStep(),
+        Krea2PrepareLatentsWithStrengthStep(),
+        Krea2RoPEInputsStep(),
+        Krea2DenoiseStep(),
+        Krea2AfterDenoiseStep(),
+    ]
+    block_names = [
+        "input",
+        "prepare_latents",
+        "set_timesteps",
+        "prepare_img2img_latents",
+        "prepare_rope_inputs",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self):
+        return "step that denoise noise into image for img2img task. It includes the denoise loop, as well as prepare the inputs (timesteps, latents, rope inputs etc.)."
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# Auto denoise step for Krea 2
+class Krea2AutoCoreDenoiseStep(AutoPipelineBlocks):
+    model_name = "krea2"
+    block_classes = [
+        Krea2InpaintCoreDenoiseStep,
+        Krea2Img2ImgCoreDenoiseStep,
+        Krea2CoreDenoiseStep,
+    ]
+    block_names = [
+        "inpaint",
+        "img2img",
+        "text2image",
+    ]
+    block_trigger_inputs = ["processed_mask_image", "image_latents", None]
+
+    @property
+    def description(self):
+        return (
+            "Core step that performs the denoising process. \n"
+            + " - `Krea2InpaintCoreDenoiseStep` (inpaint) is used when `processed_mask_image` is provided.\n"
+            + " - `Krea2Img2ImgCoreDenoiseStep` (img2img) is used when `image_latents` is provided.\n"
+            + " - `Krea2CoreDenoiseStep` (text2image) is used otherwise.\n"
+            + "This step support text-to-image, image-to-image, and inpainting tasks for Krea 2:\n"
+            + " - for image-to-image generation, you need to provide `image_latents`\n"
+            + " - for inpainting, you need to provide `processed_mask_image` and `image_latents`\n"
+            + " - for text-to-image generation, all you need to provide is prompt embeddings"
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# ====================
+# 4. DECODE
+# ====================
+
+
+# standard decode step works for most tasks except for inpaint
+# auto_docstring
+class Krea2DecodeStep(SequentialPipelineBlocks):
+    """
+    Decode step that decodes the latents to images and postprocess the generated image.
+
+      Components:
+          vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
+
+      Inputs:
+          latents (`Tensor`):
+              The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+              step.
+          output_type (`str`, *optional*, defaults to pil):
+              Output format: 'pil', 'np', 'pt'.
+
+      Outputs:
+          images (`list`):
+              Generated images. (tensor output of the vae decoder.)
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2DecoderStep(), Krea2ProcessImagesOutputStep()]
+    block_names = ["decode", "postprocess"]
+
+    @property
+    def description(self):
+        return "Decode step that decodes the latents to images and postprocess the generated image."
+
+
+# Inpaint decode step
+# auto_docstring
+class Krea2InpaintDecodeStep(SequentialPipelineBlocks):
+    """
+    Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask
+    overlay to the original image.
+
+      Components:
+          vae (`AutoencoderKLQwenImage`) image_mask_processor (`InpaintProcessor`)
+
+      Inputs:
+          latents (`Tensor`):
+              The denoised latents to decode, can be generated in the denoise step and unpacked in the after denoise
+              step.
+          output_type (`str`, *optional*, defaults to pil):
+              Output format: 'pil', 'np', 'pt'.
+          mask_overlay_kwargs (`dict`, *optional*):
+              The kwargs for the postprocess step to apply the mask overlay. generated in
+              Krea2InpaintProcessImagesInputStep.
+
+      Outputs:
+          images (`list`):
+              Generated images. (tensor output of the vae decoder.)
+    """
+
+    model_name = "krea2"
+    block_classes = [Krea2DecoderStep(), Krea2InpaintProcessImagesOutputStep()]
+    block_names = ["decode", "postprocess"]
+
+    @property
+    def description(self):
+        return "Decode step that decodes the latents to images and postprocess the generated image, optionally apply the mask overlay to the original image."
+
+
+# Auto decode step for Krea 2
+class Krea2AutoDecodeStep(AutoPipelineBlocks):
+    model_name = "krea2"
+    block_classes = [Krea2InpaintDecodeStep, Krea2DecodeStep]
+    block_names = ["inpaint_decode", "decode"]
+    block_trigger_inputs = ["mask", None]
+
+    @property
+    def description(self):
+        return (
+            "Decode step that decode the latents into images. \n"
+            " This is an auto pipeline block that works for inpaint/text2image/img2img tasks.\n"
+            + " - `Krea2InpaintDecodeStep` (inpaint_decode) is used when `mask` is provided.\n"
+            + " - `Krea2DecodeStep` (decode) is used when `mask` is not provided.\n"
+        )
+
+
+# ====================
+# 5. AUTO BLOCKS & PRESETS
+# ====================
+AUTO_BLOCKS = InsertableDict(
+    [
+        ("text_encoder", Krea2AutoTextEncoderStep()),
+        ("vae_encoder", Krea2AutoVaeEncoderStep()),
+        ("denoise", Krea2AutoCoreDenoiseStep()),
+        ("decode", Krea2AutoDecodeStep()),
+    ]
+)
+
+
+# auto_docstring
+class Krea2AutoBlocks(SequentialPipelineBlocks):
+    """
+    Auto Modular pipeline for text-to-image, image-to-image, and inpainting tasks using Krea 2.
+
+      Supported workflows:
+        - `text2image`: requires `prompt`
+        - `image2image`: requires `prompt`, `image`
+        - `inpainting`: requires `prompt`, `mask_image`, `image`
+
+      Components:
+          text_encoder (`Qwen3VLModel`): The text encoder to use tokenizer (`Qwen2Tokenizer`): The tokenizer to use
+          guider (`ClassifierFreeGuidance`) image_mask_processor (`InpaintProcessor`) vae (`AutoencoderKLQwenImage`)
+          image_processor (`VaeImageProcessor`) pachifier (`Krea2Pachifier`) scheduler
+          (`FlowMatchEulerDiscreteScheduler`) transformer (`Krea2Transformer2DModel`)
+
+      Inputs:
+          prompt (`str`, *optional*):
+              The prompt or prompts to guide image generation.
+          negative_prompt (`str`, *optional*):
+              The prompt or prompts not to guide the image generation.
+          max_sequence_length (`int`, *optional*, defaults to 512):
+              Maximum sequence length for prompt encoding.
+          mask_image (`Image`, *optional*):
+              Mask image for inpainting.
+          image (`Image | list`, *optional*):
+              Reference image(s) for denoising. Can be a single image or list of images.
+          height (`int`, *optional*):
+              The height in pixels of the generated image.
+          width (`int`, *optional*):
+              The width in pixels of the generated image.
+          padding_mask_crop (`int`, *optional*):
+              Padding for mask cropping in inpainting.
+          generator (`Generator`, *optional*):
+              Torch generator for deterministic generation.
+          num_images_per_prompt (`int`, *optional*, defaults to 1):
+              The number of images to generate per prompt.
+          prompt_embeds (`Tensor`):
+              text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          prompt_embeds_mask (`Tensor`):
+              mask for the text embeddings. Can be generated from text_encoder step.
+          negative_prompt_embeds (`Tensor`, *optional*):
+              negative text embeddings used to guide the image generation. Can be generated from text_encoder step.
+          negative_prompt_embeds_mask (`Tensor`, *optional*):
+              mask for the negative text embeddings. Can be generated from text_encoder step.
+          image_latents (`Tensor`, *optional*):
+              image latents used to guide the image generation. Can be generated from vae_encoder step.
+          processed_mask_image (`Tensor`, *optional*):
+              The processed mask image
+          latents (`Tensor`):
+              Pre-generated noisy latents for image generation.
+          num_inference_steps (`int`):
+              The number of denoising steps.
+          sigmas (`list`, *optional*):
+              Custom sigmas for the denoising process.
+          mu (`float`, *optional*):
+              Fixed timestep shift for the scheduler. Pass `1.15` for the few-step distilled (TDM/turbo) checkpoint; if
+              not provided, computed from the image sequence length (base checkpoint behavior).
+          strength (`float`, *optional*, defaults to 0.9):
+              Strength for img2img/inpainting.
+          **denoiser_input_fields (`None`, *optional*):
+              conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc.
+          output_type (`str`, *optional*, defaults to pil):
+              Output format: 'pil', 'np', 'pt'.
+          mask_overlay_kwargs (`dict`, *optional*):
+              The kwargs for the postprocess step to apply the mask overlay. generated in
+              Krea2InpaintProcessImagesInputStep.
+
+      Outputs:
+          images (`list`):
+              Generated images.
+    """
+
+    model_name = "krea2"
+
+    block_classes = AUTO_BLOCKS.values()
+    block_names = AUTO_BLOCKS.keys()
+
+    # Workflow map defines the trigger conditions for each workflow.
+    # How to define:
+    #   - Only include required inputs and trigger inputs (inputs that determine which blocks run)
+    #   - currently, only supports `True` means the workflow triggers when the input is not None
+
+    _workflow_map = {
+        "text2image": {"prompt": True},
+        "image2image": {"prompt": True, "image": True},
+        "inpainting": {"prompt": True, "mask_image": True, "image": True},
+    }
+
+    @property
+    def description(self):
+        return "Auto Modular pipeline for text-to-image, image-to-image, and inpainting tasks using Krea 2."
+
+    @property
+    def outputs(self):
+        return [OutputParam.template("images")]

modular_blocks_krea2_edit.py

@@ -0,0 +1,97 @@
+# 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.
+
+from diffusers.utils import logging
+from diffusers.modular_pipelines.modular_pipeline import SequentialPipelineBlocks
+from diffusers.modular_pipelines.modular_pipeline_utils import InsertableDict, OutputParam
+from .before_denoise import (
+    Krea2EditRoPEInputsStep,
+    Krea2PrepareLatentsStep,
+    Krea2SetTimestepsStep,
+)
+from .decoders import Krea2AfterDenoiseStep
+from .denoise import Krea2EditDenoiseStep
+from .encoders import (
+    Krea2EditReferenceLatentsStep,
+    Krea2EditTextEncoderStep,
+)
+from .inputs import Krea2TextInputsStep
+from .modular_blocks_krea2 import Krea2DecodeStep
+
+
+logger = logging.get_logger(__name__)
+
+
+# Krea 2 (reference-image edit)
+class Krea2EditCoreDenoiseStep(SequentialPipelineBlocks):
+    model_name = "krea2"
+    block_classes = [
+        Krea2TextInputsStep(),
+        Krea2PrepareLatentsStep(),
+        Krea2SetTimestepsStep(),
+        Krea2EditRoPEInputsStep(),
+        Krea2EditDenoiseStep(),
+        Krea2AfterDenoiseStep(),
+    ]
+    block_names = [
+        "input",
+        "prepare_latents",
+        "set_timesteps",
+        "prepare_rope_inputs",
+        "denoise",
+        "after_denoise",
+    ]
+
+    @property
+    def description(self):
+        return (
+            "step that denoises noise into an image for the reference-image edit task. It includes the denoise loop "
+            "(with the clean reference tokens appended each step), as well as preparing the inputs (timesteps, "
+            "latents, rope inputs etc.)."
+        )
+
+    @property
+    def outputs(self):
+        return [
+            OutputParam.template("latents"),
+        ]
+
+
+# Preset: reference-image edit. `image` carries the reference image(s); without it, use the text2image preset.
+EDIT_BLOCKS = InsertableDict(
+    [
+        ("text_encoder", Krea2EditTextEncoderStep()),
+        ("reference_latents", Krea2EditReferenceLatentsStep()),
+        ("denoise", Krea2EditCoreDenoiseStep()),
+        ("decode", Krea2DecodeStep()),
+    ]
+)
+
+
+class Krea2EditBlocks(SequentialPipelineBlocks):
+    """
+    Modular pipeline for the Krea 2 reference-image edit task.
+
+    Pass one or more reference images as `image`; they condition generation both through the Qwen3-VL text encoder
+    and as clean VAE latents appended to the transformer sequence at flow time t=0 (matching the Ostris AI-Toolkit
+    edit LoRAs). Load an edit/style LoRA into the transformer to control the effect.
+    """
+
+    model_name = "krea2"
+    block_classes = EDIT_BLOCKS.values()
+    block_names = EDIT_BLOCKS.keys()
+
+    @property
+    def description(self):
+        return "Auto Modular pipeline for the reference-image edit task using Krea 2. Requires `prompt` and `image`."

modular_pipeline.py

@@ -0,0 +1,135 @@
+# 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.
+
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.modular_pipelines.modular_pipeline import ModularPipeline
+
+
+# Copied from diffusers.modular_pipelines.qwenimage.modular_pipeline.QwenImagePachifier with QwenImage->Krea2
+class Krea2Pachifier(ConfigMixin):
+    """
+    A class to pack and unpack latents for Krea2.
+    """
+
+    config_name = "config.json"
+
+    @register_to_config
+    def __init__(self, patch_size: int = 2):
+        super().__init__()
+
+    def pack_latents(self, latents):
+        if latents.ndim != 4 and latents.ndim != 5:
+            raise ValueError(f"Latents must have 4 or 5 dimensions, but got {latents.ndim}")
+
+        if latents.ndim == 4:
+            latents = latents.unsqueeze(2)
+
+        batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width = latents.shape
+        patch_size = self.config.patch_size
+
+        if latent_height % patch_size != 0 or latent_width % patch_size != 0:
+            raise ValueError(
+                f"Latent height and width must be divisible by {patch_size}, but got {latent_height} and {latent_width}"
+            )
+
+        latents = latents.view(
+            batch_size,
+            num_channels_latents,
+            latent_height // patch_size,
+            patch_size,
+            latent_width // patch_size,
+            patch_size,
+        )
+        latents = latents.permute(
+            0, 2, 4, 1, 3, 5
+        )  # Batch_size, num_patches_height, num_patches_width, num_channels_latents, patch_size, patch_size
+        latents = latents.reshape(
+            batch_size,
+            (latent_height // patch_size) * (latent_width // patch_size),
+            num_channels_latents * patch_size * patch_size,
+        )
+
+        return latents
+
+    def unpack_latents(self, latents, height, width, vae_scale_factor=8):
+        if latents.ndim != 3:
+            raise ValueError(f"Latents must have 3 dimensions, but got {latents.ndim}")
+
+        batch_size, num_patches, channels = latents.shape
+        patch_size = self.config.patch_size
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = patch_size * (int(height) // (vae_scale_factor * patch_size))
+        width = patch_size * (int(width) // (vae_scale_factor * patch_size))
+
+        latents = latents.view(
+            batch_size,
+            height // patch_size,
+            width // patch_size,
+            channels // (patch_size * patch_size),
+            patch_size,
+            patch_size,
+        )
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (patch_size * patch_size), 1, height, width)
+
+        return latents
+
+
+class Krea2ModularPipeline(ModularPipeline):
+    """
+    A ModularPipeline for Krea 2.
+
+    > [!WARNING] > This is an experimental feature and is likely to change in the future.
+    """
+
+    default_blocks_name = "Krea2AutoBlocks"
+
+    @property
+    def default_height(self):
+        return self.default_sample_size * self.vae_scale_factor
+
+    @property
+    def default_width(self):
+        return self.default_sample_size * self.vae_scale_factor
+
+    @property
+    def default_sample_size(self):
+        return 128
+
+    @property
+    def vae_scale_factor(self):
+        vae_scale_factor = 8
+        if hasattr(self, "vae") and self.vae is not None:
+            vae_scale_factor = 2 ** len(self.vae.temperal_downsample)
+        return vae_scale_factor
+
+    @property
+    def num_channels_latents(self):
+        num_channels_latents = 16
+        if hasattr(self, "transformer") and self.transformer is not None:
+            num_channels_latents = self.transformer.config.in_channels // 4
+        return num_channels_latents
+
+    @property
+    def requires_unconditional_embeds(self):
+        requires_unconditional_embeds = False
+
+        if hasattr(self, "guider") and self.guider is not None:
+            requires_unconditional_embeds = self.guider._enabled and self.guider.num_conditions > 1
+
+        return requires_unconditional_embeds

transformer_krea2.py

@@ -0,0 +1,566 @@
+# 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.
+
+import inspect
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.utils import logging
+from diffusers.utils.torch_utils import maybe_adjust_dtype_for_device
+from diffusers.models.attention import AttentionMixin, AttentionModuleMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.embeddings import apply_rotary_emb, get_1d_rotary_pos_embed
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+class Krea2RMSNorm(nn.Module):
+    """RMSNorm with a zero-centered scale: the effective multiplier is `1 + weight`, matching the Krea 2 checkpoint
+    format. The activations are upcast so the normalization runs in float32; the scale weight is kept in float32 by
+    the model's `_keep_in_fp32_modules`."""
+
+    def __init__(self, dim: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.dim = dim
+        self.eps = eps
+        self.weight = nn.Parameter(torch.zeros(dim))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        dtype = hidden_states.dtype
+        hidden_states = F.rms_norm(hidden_states.float(), (self.dim,), weight=self.weight + 1.0, eps=self.eps)
+        return hidden_states.to(dtype)
+
+
+class Krea2AttnProcessor:
+    _attention_backend = None
+    _parallel_config = None
+
+    def __call__(
+        self,
+        attn: "Krea2Attention",
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        query = attn.to_q(hidden_states).unflatten(-1, (attn.num_heads, attn.head_dim))
+        key = attn.to_k(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
+        value = attn.to_v(hidden_states).unflatten(-1, (attn.num_kv_heads, attn.head_dim))
+        gate = attn.to_gate(hidden_states)
+
+        query = attn.norm_q(query)
+        key = attn.norm_k(key)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
+            key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)
+
+        hidden_states = dispatch_attention_fn(
+            query,
+            key,
+            value,
+            attn_mask=attention_mask,
+            enable_gqa=attn.num_heads != attn.num_kv_heads,
+            backend=self._attention_backend,
+            parallel_config=self._parallel_config,
+        )
+        hidden_states = hidden_states.flatten(2, 3)
+        hidden_states = hidden_states * torch.sigmoid(gate)
+        return attn.to_out[0](hidden_states)
+
+
+class Krea2Attention(nn.Module, AttentionModuleMixin):
+    """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings and a sigmoid output gate."""
+
+    _default_processor_cls = Krea2AttnProcessor
+    _available_processors = [Krea2AttnProcessor]
+
+    def __init__(
+        self, hidden_size: int, num_heads: int, num_kv_heads: int | None = None, eps: float = 1e-5, processor=None
+    ) -> None:
+        super().__init__()
+        if hidden_size % num_heads != 0:
+            raise ValueError(f"hidden_size={hidden_size} must be divisible by num_heads={num_heads}")
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+        self.head_dim = hidden_size // num_heads
+        self.use_bias = False
+
+        self.to_q = nn.Linear(hidden_size, self.head_dim * self.num_heads, bias=False)
+        self.to_k = nn.Linear(hidden_size, self.head_dim * self.num_kv_heads, bias=False)
+        self.to_v = nn.Linear(hidden_size, self.head_dim * self.num_kv_heads, bias=False)
+        self.to_gate = nn.Linear(hidden_size, hidden_size, bias=False)
+        self.norm_q = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.norm_k = Krea2RMSNorm(self.head_dim, eps=eps)
+        self.to_out = nn.ModuleList([nn.Linear(hidden_size, hidden_size, bias=False), nn.Dropout(0.0)])
+
+        if processor is None:
+            processor = self._default_processor_cls()
+        self.set_processor(processor)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+        unused_kwargs = [k for k in kwargs if k not in attn_parameters]
+        if len(unused_kwargs) > 0:
+            logger.warning(
+                f"attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+            )
+        kwargs = {k: w for k, w in kwargs.items() if k in attn_parameters}
+        return self.processor(self, hidden_states, attention_mask, image_rotary_emb, **kwargs)
+
+
+class Krea2SwiGLU(nn.Module):
+    """SwiGLU feed-forward network."""
+
+    def __init__(self, dim: int, hidden_dim: int) -> None:
+        super().__init__()
+        self.gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.up = nn.Linear(dim, hidden_dim, bias=False)
+        self.down = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.down(F.silu(self.gate(hidden_states)) * self.up(hidden_states))
+
+
+class Krea2TextFusionBlock(nn.Module):
+    """Pre-norm transformer block (no rotary embeddings, no time modulation) used by the text fusion stage."""
+
+    def __init__(self, dim: int, num_heads: int, num_kv_heads: int, intermediate_size: int, eps: float) -> None:
+        super().__init__()
+        self.norm1 = Krea2RMSNorm(dim, eps=eps)
+        self.norm2 = Krea2RMSNorm(dim, eps=eps)
+        self.attn = Krea2Attention(dim, num_heads, num_kv_heads, eps=eps)
+        self.ff = Krea2SwiGLU(dim, intermediate_size)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        hidden_states = hidden_states + self.attn(self.norm1(hidden_states), attention_mask=attention_mask)
+        hidden_states = hidden_states + self.ff(self.norm2(hidden_states))
+        return hidden_states
+
+
+class Krea2TextFusion(nn.Module):
+    """Fuses the stack of tapped text-encoder hidden states into a single sequence of text features.
+
+    Two `layerwise_blocks` attend across the `num_text_layers` axis independently for every token, a linear
+    `projector` collapses that axis, and two `refiner_blocks` attend across the token sequence.
+    """
+
+    def __init__(
+        self,
+        num_text_layers: int,
+        dim: int,
+        num_heads: int,
+        num_kv_heads: int,
+        intermediate_size: int,
+        num_layerwise_blocks: int,
+        num_refiner_blocks: int,
+        eps: float,
+    ) -> None:
+        super().__init__()
+        self.layerwise_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(dim, num_heads, num_kv_heads, intermediate_size, eps)
+                for _ in range(num_layerwise_blocks)
+            ]
+        )
+        self.projector = nn.Linear(num_text_layers, 1, bias=False)
+        self.refiner_blocks = nn.ModuleList(
+            [
+                Krea2TextFusionBlock(dim, num_heads, num_kv_heads, intermediate_size, eps)
+                for _ in range(num_refiner_blocks)
+            ]
+        )
+
+    def forward(self, encoder_hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None) -> torch.Tensor:
+        batch_size, seq_len, num_text_layers, dim = encoder_hidden_states.shape
+
+        hidden_states = encoder_hidden_states.reshape(batch_size * seq_len, num_text_layers, dim)
+        for block in self.layerwise_blocks:
+            hidden_states = block(hidden_states.contiguous())
+
+        hidden_states = hidden_states.reshape(batch_size, seq_len, num_text_layers, dim).permute(0, 1, 3, 2)
+        hidden_states = self.projector(hidden_states).squeeze(-1)
+
+        for block in self.refiner_blocks:
+            hidden_states = block(hidden_states, attention_mask=attention_mask)
+
+        return hidden_states
+
+
+class Krea2TransformerBlock(nn.Module):
+    def __init__(
+        self, hidden_size: int, intermediate_size: int, num_heads: int, num_kv_heads: int, norm_eps: float
+    ) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(6, hidden_size))
+        self.norm1 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.norm2 = Krea2RMSNorm(hidden_size, eps=norm_eps)
+        self.attn = Krea2Attention(hidden_size, num_heads, num_kv_heads, eps=norm_eps)
+        self.ff = Krea2SwiGLU(hidden_size, intermediate_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        temb: torch.Tensor | tuple[torch.Tensor, torch.Tensor, int],
+        image_rotary_emb: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # temb: (B, 1, 6 * hidden_size), shared across all blocks; each block only learns an additive table.
+        # For reference-image ("edit") conditioning, temb is instead a tuple (temb, ref_temb, split): the leading
+        # `split` tokens (text + noisy image) are modulated with the real timestep while the trailing tokens (clean
+        # reference tokens) use the t=0 embedding `ref_temb`.
+        if isinstance(temb, tuple):
+            temb, ref_temb, split = temb
+            m = (temb.unflatten(-1, (6, -1)) + self.scale_shift_table).unbind(-2)
+            r = (ref_temb.unflatten(-1, (6, -1)) + self.scale_shift_table).unbind(-2)
+
+            def modulate(h, scale_idx, shift_idx):
+                return torch.cat(
+                    (
+                        (1.0 + m[scale_idx]) * h[:, :split] + m[shift_idx],
+                        (1.0 + r[scale_idx]) * h[:, split:] + r[shift_idx],
+                    ),
+                    dim=1,
+                )
+
+            def gate(h, gate_idx):
+                return torch.cat((m[gate_idx] * h[:, :split], r[gate_idx] * h[:, split:]), dim=1)
+
+            attn_out = self.attn(
+                modulate(self.norm1(hidden_states), 0, 1),
+                attention_mask=attention_mask,
+                image_rotary_emb=image_rotary_emb,
+            )
+            hidden_states = hidden_states + gate(attn_out, 2)
+            ff_out = self.ff(modulate(self.norm2(hidden_states), 3, 4))
+            hidden_states = hidden_states + gate(ff_out, 5)
+            return hidden_states
+
+        modulation = temb.unflatten(-1, (6, -1)) + self.scale_shift_table
+        prescale, preshift, pregate, postscale, postshift, postgate = modulation.unbind(-2)
+
+        attn_out = self.attn(
+            (1.0 + prescale) * self.norm1(hidden_states) + preshift,
+            attention_mask=attention_mask,
+            image_rotary_emb=image_rotary_emb,
+        )
+        hidden_states = hidden_states + pregate * attn_out
+        ff_out = self.ff((1.0 + postscale) * self.norm2(hidden_states) + postshift)
+        hidden_states = hidden_states + postgate * ff_out
+        return hidden_states
+
+
+class Krea2TimestepEmbedding(nn.Module):
+    """Sinusoidal flow-time embedding (cos-first, input scaled by 1000) followed by a two-layer MLP.
+
+    Keeps the sequence dimension at size 1 so the per-block modulations broadcast over tokens.
+    """
+
+    def __init__(self, embed_dim: int, hidden_size: int) -> None:
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.linear_1 = nn.Linear(embed_dim, hidden_size, bias=True)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, timestep: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+        half = self.embed_dim // 2
+        freqs = torch.exp(-math.log(1e4) * torch.arange(half, dtype=torch.float32, device=timestep.device) / half)
+        args = (timestep.float() * 1e3)[:, None, None] * freqs
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1).to(dtype)
+        return self.linear_2(F.gelu(self.linear_1(emb), approximate="tanh"))
+
+
+class Krea2TextProjection(nn.Module):
+    """Projects the fused text features into the transformer width."""
+
+    def __init__(self, text_dim: int, hidden_size: int, eps: float) -> None:
+        super().__init__()
+        self.norm = Krea2RMSNorm(text_dim, eps=eps)
+        self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True)
+        self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.linear_1(self.norm(hidden_states))
+        return self.linear_2(F.gelu(hidden_states, approximate="tanh"))
+
+
+class Krea2FinalLayer(nn.Module):
+    """Final adaptive RMSNorm and output projection. Kept as one module (and in `_no_split_modules`) so the learned
+    modulation table, norm and projection stay co-located under device-mapped inference."""
+
+    def __init__(self, hidden_size: int, out_channels: int, eps: float) -> None:
+        super().__init__()
+        self.scale_shift_table = nn.Parameter(torch.zeros(2, hidden_size))
+        self.norm = Krea2RMSNorm(hidden_size, eps=eps)
+        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, hidden_states: torch.Tensor, temb: torch.Tensor) -> torch.Tensor:
+        modulation = temb + self.scale_shift_table
+        scale, shift = modulation.chunk(2, dim=1)
+        hidden_states = (1.0 + scale) * self.norm(hidden_states) + shift
+        return self.linear(hidden_states)
+
+
+# Copied from diffusers.models.transformers.transformer_flux.FluxPosEmbed with FluxPosEmbed->Krea2RotaryPosEmbed
+class Krea2RotaryPosEmbed(nn.Module):
+    # modified from https://github.com/black-forest-labs/flux/blob/c00d7c60b085fce8058b9df845e036090873f2ce/src/flux/modules/layers.py#L11
+    def __init__(self, theta: int, axes_dim: list[int]):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> torch.Tensor:
+        n_axes = ids.shape[-1]
+        cos_out = []
+        sin_out = []
+        pos = ids.float()
+        freqs_dtype = maybe_adjust_dtype_for_device(torch.float64, ids.device)
+        for i in range(n_axes):
+            cos, sin = get_1d_rotary_pos_embed(
+                self.axes_dim[i],
+                pos[:, i],
+                theta=self.theta,
+                repeat_interleave_real=True,
+                use_real=True,
+                freqs_dtype=freqs_dtype,
+            )
+            cos_out.append(cos)
+            sin_out.append(sin)
+        freqs_cos = torch.cat(cos_out, dim=-1).to(ids.device)
+        freqs_sin = torch.cat(sin_out, dim=-1).to(ids.device)
+        return freqs_cos, freqs_sin
+
+
+class Krea2Transformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin):
+    r"""
+    The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline.
+
+    Text conditioning enters as a stack of hidden states tapped from several layers of a multimodal text encoder. A
+    small text-fusion transformer collapses the layer axis and refines the token sequence; the result is concatenated
+    with the patchified image latents into a single `[text, image]` sequence processed by the transformer blocks. The
+    timestep conditions every block through one shared modulation vector plus per-block learned tables.
+
+    Args:
+        in_channels (`int`, defaults to 64):
+            Latent channel count after patchification (`vae_channels * patch_size ** 2`).
+        num_layers (`int`, defaults to 28):
+            Number of transformer blocks.
+        attention_head_dim (`int`, defaults to 128):
+            Dimension of each attention head; the total hidden size is `attention_head_dim * num_attention_heads`.
+        num_attention_heads (`int`, defaults to 48):
+            Number of query heads.
+        num_key_value_heads (`int`, defaults to 12):
+            Number of key/value heads for grouped-query attention.
+        intermediate_size (`int`, defaults to 16384):
+            Feed-forward hidden size of the SwiGLU MLP inside each block.
+        timestep_embed_dim (`int`, defaults to 256):
+            Width of the sinusoidal timestep embedding before its MLP.
+        text_hidden_dim (`int`, defaults to 2560):
+            Hidden size of the text encoder whose hidden states are consumed.
+        num_text_layers (`int`, defaults to 12):
+            Number of tapped text-encoder hidden states stacked per token.
+        text_num_attention_heads (`int`, defaults to 20):
+            Number of query heads in the text fusion blocks.
+        text_num_key_value_heads (`int`, defaults to 20):
+            Number of key/value heads in the text fusion blocks.
+        text_intermediate_size (`int`, defaults to 6912):
+            Feed-forward hidden size of the SwiGLU MLP inside the text fusion blocks.
+        num_layerwise_text_blocks (`int`, defaults to 2):
+            Number of text fusion blocks applied across the tapped-layer axis (per token).
+        num_refiner_text_blocks (`int`, defaults to 2):
+            Number of text fusion blocks applied across the token sequence.
+        axes_dims_rope (`tuple[int, int, int]`, defaults to `(32, 48, 48)`):
+            Head-dim split across the (t, h, w) rotary position axes.
+        rope_theta (`float`, defaults to 1000.0):
+            Base used by the rotary position embedding.
+        norm_eps (`float`, defaults to 1e-5):
+            Epsilon used by all RMSNorm modules.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["Krea2TransformerBlock", "Krea2TextFusionBlock", "Krea2FinalLayer"]
+    _repeated_blocks = ["Krea2TransformerBlock"]
+    _keep_in_fp32_modules = ["norm", "norm1", "norm2", "norm_q", "norm_k"]
+    _skip_layerwise_casting_patterns = ["time_embed", "norm"]
+
+    @register_to_config
+    def __init__(
+        self,
+        in_channels: int = 64,
+        num_layers: int = 28,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 48,
+        num_key_value_heads: int = 12,
+        intermediate_size: int = 16384,
+        timestep_embed_dim: int = 256,
+        text_hidden_dim: int = 2560,
+        num_text_layers: int = 12,
+        text_num_attention_heads: int = 20,
+        text_num_key_value_heads: int = 20,
+        text_intermediate_size: int = 6912,
+        num_layerwise_text_blocks: int = 2,
+        num_refiner_text_blocks: int = 2,
+        axes_dims_rope: tuple[int, int, int] = (32, 48, 48),
+        rope_theta: float = 1000.0,
+        norm_eps: float = 1e-5,
+    ) -> None:
+        super().__init__()
+
+        hidden_size = attention_head_dim * num_attention_heads
+        if sum(axes_dims_rope) != attention_head_dim:
+            raise ValueError(
+                f"sum(axes_dims_rope)={sum(axes_dims_rope)} must equal attention_head_dim={attention_head_dim}"
+            )
+
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.hidden_size = hidden_size
+        self.gradient_checkpointing = False
+
+        self.img_in = nn.Linear(in_channels, hidden_size, bias=True)
+        self.time_embed = Krea2TimestepEmbedding(timestep_embed_dim, hidden_size)
+        self.time_mod_proj = nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+        self.text_fusion = Krea2TextFusion(
+            num_text_layers=num_text_layers,
+            dim=text_hidden_dim,
+            num_heads=text_num_attention_heads,
+            num_kv_heads=text_num_key_value_heads,
+            intermediate_size=text_intermediate_size,
+            num_layerwise_blocks=num_layerwise_text_blocks,
+            num_refiner_blocks=num_refiner_text_blocks,
+            eps=norm_eps,
+        )
+        self.txt_in = Krea2TextProjection(text_hidden_dim, hidden_size, eps=norm_eps)
+        self.rotary_emb = Krea2RotaryPosEmbed(theta=rope_theta, axes_dim=list(axes_dims_rope))
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                Krea2TransformerBlock(
+                    hidden_size=hidden_size,
+                    intermediate_size=intermediate_size,
+                    num_heads=num_attention_heads,
+                    num_kv_heads=num_key_value_heads,
+                    norm_eps=norm_eps,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.final_layer = Krea2FinalLayer(hidden_size, out_channels=in_channels, eps=norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        position_ids: torch.Tensor,
+        encoder_attention_mask: torch.Tensor | None = None,
+        ref_seq_len: int = 0,
+        return_dict: bool = True,
+    ) -> Transformer2DModelOutput | tuple[torch.Tensor]:
+        r"""
+        Predict the flow-matching velocity for the (noisy) image tokens.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_seq_len + ref_seq_len, in_channels)`):
+                Packed (patchified) noisy image latents, with any packed clean reference latents appended at the end.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_seq_len, num_text_layers, text_hidden_dim)`):
+                Stack of tapped text-encoder hidden states per token.
+            timestep (`torch.Tensor` of shape `(batch_size,)`):
+                Flow-matching time in `[0, 1]` (1 is pure noise, 0 is clean data).
+            position_ids (`torch.Tensor` of shape `(text_seq_len + image_seq_len + ref_seq_len, 3)`):
+                `(t, h, w)` rotary coordinates for the combined sequence. Text rows are all-zero; image rows hold the
+                latent-grid coordinates; the i-th reference image sits on frame axis `i + 1` with its own grid.
+            encoder_attention_mask (`torch.Tensor` of shape `(batch_size, text_seq_len)`, *optional*):
+                Boolean mask marking valid text tokens. Pass `None` when every text token is valid.
+            ref_seq_len (`int`, *optional*, defaults to 0):
+                Number of trailing reference (edit) tokens in `hidden_states`. They are modulated with the t=0
+                embedding and excluded from the returned velocity.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain tuple.
+
+        Returns:
+            [`~models.modeling_outputs.Transformer2DModelOutput`] or a `tuple` whose first element is the velocity
+            tensor of shape `(batch_size, image_seq_len, in_channels)`.
+        """
+        if position_ids.ndim != 2 or position_ids.shape[-1] != 3:
+            raise ValueError(f"`position_ids` must have shape (sequence_length, 3), got {tuple(position_ids.shape)}.")
+
+        batch_size, image_seq_len, _ = hidden_states.shape  # image_seq_len includes any trailing reference tokens
+        text_seq_len = encoder_hidden_states.shape[1]
+
+        temb = self.time_embed(timestep, dtype=hidden_states.dtype)
+        temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh"))
+
+        # Clean reference tokens are conditioned at flow time t=0; the text + noisy image tokens keep the real
+        # timestep. The blocks then receive a (temb, ref_temb, split) tuple that modulates the two spans separately.
+        block_temb = temb_mod
+        if ref_seq_len > 0:
+            temb_zero = self.time_embed(torch.zeros_like(timestep), dtype=hidden_states.dtype)
+            ref_temb_mod = self.time_mod_proj(F.gelu(temb_zero, approximate="tanh"))
+            block_temb = (temb_mod, ref_temb_mod, text_seq_len + image_seq_len - ref_seq_len)
+
+        text_attention_mask = None
+        attention_mask = None
+        if encoder_attention_mask is not None:
+            # Key-padding masks of shape (B, 1, 1, L): padded text tokens are excluded as attention keys everywhere;
+            # their own (garbage) lanes are never read back and are dropped at the output slice.
+            text_attention_mask = encoder_attention_mask[:, None, None, :]
+            image_mask = encoder_attention_mask.new_ones((batch_size, image_seq_len))
+            attention_mask = torch.cat([encoder_attention_mask, image_mask], dim=1)[:, None, None, :]
+
+        encoder_hidden_states = self.text_fusion(encoder_hidden_states, attention_mask=text_attention_mask)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        hidden_states = self.img_in(hidden_states)
+        hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+        image_rotary_emb = self.rotary_emb(position_ids)
+
+        for block in self.transformer_blocks:
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                hidden_states = self._gradient_checkpointing_func(
+                    block, hidden_states, block_temb, image_rotary_emb, attention_mask
+                )
+            else:
+                hidden_states = block(hidden_states, block_temb, image_rotary_emb, attention_mask)
+
+        # Keep only the noisy image tokens: drop the leading text tokens and any trailing reference (edit) tokens.
+        hidden_states = hidden_states[:, text_seq_len : text_seq_len + image_seq_len - ref_seq_len]
+        output = self.final_layer(hidden_states, temb)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+
+# Register into the diffusers namespace so `ModularPipeline` component loading can resolve the
+# transformer referenced as ["diffusers", "Krea2Transformer2DModel"] in the base repo's model_index.json.
+import diffusers as _diffusers  # noqa: E402
+
+_diffusers.Krea2Transformer2DModel = Krea2Transformer2DModel