Upload private Krea 2 Outpaint release candidate

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+LICENSE.pdf filter=lfs diff=lfs merge=lfs -text

LICENSE.pdf

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b82a2805162bde714a4eb27b9063c4fc3345d08a30be055134a6160e5430ba74
+size 137711

NOTICE

@@ -0,0 +1,15 @@
+Krea 2 Registered Outpaint
+
+The model weights in this repository are derived from Krea 2 and are subject to
+the Krea 2 Community License included as LICENSE.pdf.
+
+The included pipeline is derived from ostris/Krea2OstrisEdit and adds registered
+reference placement through destination-relative rotary coordinates. Pipeline
+and helper code are distributed under the Apache License 2.0.
+
+Krea 2: https://huggingface.co/krea/Krea-2-Turbo
+Upstream edit pipeline: https://huggingface.co/ostris/Krea2OstrisEdit
+
+This is an unofficial community release and is not endorsed by Krea or Ostris.
+Training data is not included.
+

PIPELINE_LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   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.

README.md

@@ -0,0 +1,145 @@
+---
+license: other
+license_name: krea-2-community-license
+license_link: https://krea.ai/krea-2-licensing
+base_model: krea/Krea-2-Turbo
+pipeline_tag: image-to-image
+tags:
+  - krea-2
+  - lora
+  - outpainting
+  - image-editing
+  - diffusers
+---
+
+# Krea 2 Registered Outpaint
+
+> Draft release candidate. This repository is private while its documentation,
+> license files, and clean-download inference path are reviewed.
+
+This rank-32 LoRA extends an image into a larger canvas while preserving its
+known pixels and registering the source image at an explicit canvas location.
+It was trained against Krea 2 Raw and is intended for distilled 8-step inference
+with Krea 2 Turbo.
+
+Unlike a plain image-reference adapter, the source latent tokens receive rotary
+coordinates derived from their destination bounding box. The denoiser can
+therefore reason about where the known image belongs in the target canvas. The
+included pipeline is based on
+[ostris/Krea2OstrisEdit](https://huggingface.co/ostris/Krea2OstrisEdit) and adds
+the registered `reference_placements` input used by this adapter.
+
+## Files
+
+```text
+krea2_outpaint_rank32.safetensors  LoRA weights
+pipeline.py                        Krea 2 edit pipeline with registered references
+outpaint.py                        Canvas preparation and seam compositing helpers
+example.py                         Minimal local example
+LICENSE.pdf                        Krea 2 Community License
+PIPELINE_LICENSE                   Apache-2.0 license for pipeline/helper code
+NOTICE                             attribution and derivative-work notice
+SHA256SUMS                         release artifact hashes
+```
+
+## Input contract
+
+The model takes:
+
+- one RGB source image;
+- an output canvas whose width and height are multiples of 16;
+- a source bounding box `(x0, y0, x1, y1)` in output pixels;
+- a prompt describing the complete output image.
+
+The bounding box must stay inside the canvas and preserve the source aspect
+ratio. The helper resizes the source to that exact box. Conditioning is encoded
+at a maximum edge of 384 pixels; the original resized pixels are restored after
+generation with a 32-pixel inward feather.
+
+For a single pass, the source box should span the full canvas width or height.
+`outpaint.py` also provides a two-pass plan for arbitrary interior boxes. It
+chooses the smaller intermediate blank region first.
+
+## Recommended settings
+
+| Setting | Value |
+|---|---:|
+| Base | `krea/Krea-2-Turbo` |
+| Steps | 8 |
+| Guidance scale | 0.0 |
+| LoRA scale | 1.0 |
+| Reference maximum edge | 384 px |
+| Seam feather | 32 px |
+| Reference VLM encoding | disabled |
+| Reference K/V cache | enabled |
+
+The adapter is an AI-Toolkit LoRA and uses isolated reference attention, so
+`kv_cache=True` and `encode_reference_in_prompt=False` are part of the model
+contract rather than optional speed tweaks.
+
+## Usage
+
+Install a current PyTorch build for your GPU, then install the dependencies:
+
+```bash
+pip install diffusers transformers accelerate safetensors huggingface_hub pillow
+```
+
+Run the included example after downloading this repository:
+
+```bash
+python example.py \
+  --source input.png \
+  --output output.png \
+  --width 1024 \
+  --height 1536 \
+  --bbox 0 0 1024 1024 \
+  --prompt "a polished digital illustration, full composition"
+```
+
+The example defaults to BF16. Quantized Krea 2 runtimes may also load the LoRA,
+but quantization setup is runtime-specific and is not bundled here.
+
+## How it works
+
+The noisy target follows normal Krea 2 flow-matching inference. The clean source
+is VAE-encoded as a separate token sequence. For a normalized source box
+`(x0, y0, x1, y1)`, source token coordinates are mapped into the target latent
+grid before rotary position encoding. Reference K/V is computed once and reused
+through all denoising steps. Finally, exact source pixels are composited back
+into the known region with a short inward feather to retain detail while hiding
+small decode differences at the boundary.
+
+## Limitations
+
+- Very large missing regions can introduce a second subject or unrelated scene.
+- Interior placement uses two generations and can accumulate changes between
+  passes.
+- The model handles rectangular outpainting, not arbitrary inpainting masks.
+- Prompt wording and source composition affect how aggressively content is
+  extended.
+- The 384-pixel conditioning budget prioritizes speed and structure over exact
+  recovery of tiny source details; final compositing preserves the known pixels.
+- BF16 clean-download inference remains a release-blocking verification item for
+  this draft.
+
+## Training summary
+
+The adapter was trained at rank/alpha 32/32 on rectangular crops registered to
+their location in complete target images. Targets were at most 1024 pixels on
+their longest edge and source conditions at most 384 pixels. The final candidate
+is the 12,000-step checkpoint. Training data and training infrastructure are not
+distributed.
+
+In the tested INT8 ConvRot Krea 2 Turbo runtime on an RTX 5090, one 8-step pass
+took approximately 3.6-4.2 seconds at the evaluated native resolutions, with
+prompt encoding around 0.1 seconds. These measurements are implementation- and
+hardware-specific and are not BF16 benchmark claims.
+
+## License
+
+The LoRA weights are a derivative of Krea 2 and are released under the Krea 2
+Community License in `LICENSE.pdf`. Pipeline and helper code are distributed
+under Apache-2.0 in `PIPELINE_LICENSE`. This is an unofficial community release
+and is not endorsed by Krea or Ostris.
+

SHA256SUMS

@@ -0,0 +1,8 @@
+b82a2805162bde714a4eb27b9063c4fc3345d08a30be055134a6160e5430ba74  LICENSE.pdf
+5d23d01efae1fde8eccd4f00215386860c89ef5071085f1e7b311bda1baafeb1  NOTICE
+cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30  PIPELINE_LICENSE
+f76772b7e9be831d8b761a3394c1731182ac9b9a4dc8a628376edca3bfbde117  README.md
+39a7e308656302da1cc2c26130ff512f8931197033743913d7bfaecac2c6ff15  example.py
+1de7c106ae978de21433152bdb4ba6b09118b33be81b7ee4851420f0013a7e76  krea2_outpaint_rank32.safetensors
+cb200ee7b4d12eb5cc1f663a661b7b7902bc955dc3d6286447a2ee77fbedc7d7  outpaint.py
+877de17e0a9f7e42700d345b84a2ae9a3a3e559fdd5a69d4fb66a1db35f4449f  pipeline.py

example.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
+
+from outpaint import composite, prepare_source
+
+
+REPO_ID = "yijunwang2/krea2-outpaint"
+WEIGHT_NAME = "krea2_outpaint_rank32.safetensors"
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Run registered Krea 2 outpainting")
+    parser.add_argument("--source", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--width", type=int, required=True)
+    parser.add_argument("--height", type=int, required=True)
+    parser.add_argument("--bbox", type=int, nargs=4, metavar=("X0", "Y0", "X1", "Y1"), required=True)
+    parser.add_argument("--prompt", required=True)
+    parser.add_argument("--steps", type=int, default=8)
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--lora-scale", type=float, default=1.0)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    prepared = prepare_source(
+        Image.open(args.source),
+        (args.width, args.height),
+        tuple(args.bbox),
+    )
+
+    pipe = DiffusionPipeline.from_pretrained(
+        "krea/Krea-2-Turbo",
+        custom_pipeline=REPO_ID,
+        torch_dtype=torch.bfloat16,
+    ).to("cuda")
+    pipe.load_lora_weights(REPO_ID, weight_name=WEIGHT_NAME, adapter_name="outpaint")
+    pipe.set_adapters(["outpaint"], adapter_weights=[args.lora_scale])
+
+    generator = torch.Generator(device="cuda").manual_seed(args.seed)
+    generated = pipe(
+        prompt=args.prompt,
+        image=prepared.condition,
+        width=args.width,
+        height=args.height,
+        num_inference_steps=args.steps,
+        guidance_scale=0.0,
+        generator=generator,
+        reference_max_pixels=384 * 384,
+        reference_placements=[{"bbox_normalized": prepared.bbox_normalized}],
+        encode_reference_in_prompt=False,
+        kv_cache=True,
+    ).images[0]
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    composite(generated, prepared).save(args.output)
+
+
+if __name__ == "__main__":
+    main()
+

krea2_outpaint_rank32.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1de7c106ae978de21433152bdb4ba6b09118b33be81b7ee4851420f0013a7e76
+size 228587736

outpaint.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+import math
+
+import numpy as np
+from PIL import Image
+
+
+SOURCE_MAX_EDGE = 384
+SEAM_PX = 32
+
+
+@dataclass(frozen=True)
+class RegisteredSource:
+    condition: Image.Image
+    placed_source: Image.Image
+    canvas_size: tuple[int, int]
+    bbox: tuple[int, int, int, int]
+    seam_px: int = SEAM_PX
+
+    @property
+    def bbox_normalized(self) -> list[float]:
+        width, height = self.canvas_size
+        x0, y0, x1, y1 = self.bbox
+        return [x0 / width, y0 / height, x1 / width, y1 / height]
+
+
+@dataclass(frozen=True)
+class PassPlan:
+    axis: str
+    intermediate_size: tuple[int, int] | None
+    first_bbox: tuple[int, int, int, int]
+    second_bbox: tuple[int, int, int, int] | None
+
+    @property
+    def pass_count(self) -> int:
+        return 1 if self.intermediate_size is None else 2
+
+
+def _resize_max_edge(image: Image.Image, max_edge: int) -> Image.Image:
+    if max(image.size) <= max_edge:
+        return image.copy()
+    scale = max_edge / max(image.size)
+    size = (max(1, round(image.width * scale)), max(1, round(image.height * scale)))
+    return image.resize(size, Image.Resampling.LANCZOS)
+
+
+def prepare_source(
+    source: Image.Image,
+    canvas_size: tuple[int, int],
+    bbox: tuple[int, int, int, int],
+    *,
+    source_max_edge: int = SOURCE_MAX_EDGE,
+    seam_px: int = SEAM_PX,
+) -> RegisteredSource:
+    width, height = canvas_size
+    x0, y0, x1, y1 = bbox
+    if width < 16 or height < 16 or width % 16 or height % 16:
+        raise ValueError("Canvas dimensions must be positive multiples of 16")
+    if not (0 <= x0 < x1 <= width and 0 <= y0 < y1 <= height):
+        raise ValueError(f"Source bbox is outside the canvas: {bbox}")
+
+    source = source.convert("RGB")
+    box_width, box_height = x1 - x0, y1 - y0
+    source_ratio = source.width / source.height
+    box_ratio = box_width / box_height
+    tolerance = max(0.025, 2.0 / min(box_width, box_height))
+    if abs(box_ratio / source_ratio - 1.0) > tolerance:
+        raise ValueError("Source bbox must preserve the source image aspect ratio")
+
+    placed = source.resize((box_width, box_height), Image.Resampling.LANCZOS)
+    return RegisteredSource(
+        condition=_resize_max_edge(placed, source_max_edge),
+        placed_source=placed,
+        canvas_size=canvas_size,
+        bbox=bbox,
+        seam_px=seam_px,
+    )
+
+
+def _align_up(value: int, alignment: int = 16) -> int:
+    return int(math.ceil(value / alignment) * alignment)
+
+
+def plan_passes(prepared: RegisteredSource) -> PassPlan:
+    width, height = prepared.canvas_size
+    x0, y0, x1, y1 = prepared.bbox
+    source_width, source_height = x1 - x0, y1 - y0
+    if source_width == width or source_height == height:
+        return PassPlan("direct", None, prepared.bbox, None)
+
+    candidates: list[tuple[int, PassPlan]] = []
+    intermediate_height = _align_up(source_height)
+    if intermediate_height < height:
+        intermediate_y = max(
+            0,
+            min(height - intermediate_height, y0 - (intermediate_height - source_height) // 2),
+        )
+        local_y = y0 - intermediate_y
+        candidates.append(
+            (
+                width * intermediate_height,
+                PassPlan(
+                    "horizontal_first",
+                    (width, intermediate_height),
+                    (x0, local_y, x1, local_y + source_height),
+                    (0, intermediate_y, width, intermediate_y + intermediate_height),
+                ),
+            )
+        )
+
+    intermediate_width = _align_up(source_width)
+    if intermediate_width < width:
+        intermediate_x = max(
+            0,
+            min(width - intermediate_width, x0 - (intermediate_width - source_width) // 2),
+        )
+        local_x = x0 - intermediate_x
+        candidates.append(
+            (
+                intermediate_width * height,
+                PassPlan(
+                    "vertical_first",
+                    (intermediate_width, height),
+                    (local_x, y0, local_x + source_width, y1),
+                    (intermediate_x, 0, intermediate_x + intermediate_width, height),
+                ),
+            )
+        )
+
+    if not candidates:
+        return PassPlan("direct", None, prepared.bbox, None)
+    return min(candidates, key=lambda item: (item[0], item[1].axis))[1]
+
+
+def composite(generated: Image.Image, prepared: RegisteredSource) -> Image.Image:
+    generated = generated.convert("RGB")
+    if generated.size != prepared.canvas_size:
+        raise ValueError("Generated image size does not match the canvas")
+
+    width, height = prepared.placed_source.size
+    yy, xx = np.mgrid[:height, :width]
+    edge_distance = np.minimum.reduce((xx, yy, width - 1 - xx, height - 1 - yy))
+    alpha = np.clip(edge_distance / max(1, prepared.seam_px), 0.0, 1.0)
+    alpha_image = Image.fromarray((alpha * 255).astype(np.uint8), mode="L")
+    result = generated.copy()
+    result.paste(prepared.placed_source, prepared.bbox[:2], alpha_image)
+    return result
+

pipeline.py

@@ -0,0 +1,1522 @@
+# Copyright 2026 Ostris, LLC. All rights reserved.
+#
+# Portions of the Krea2Transformer2DModel implementation are adapted from
+# huggingface/diffusers (Apache License, Version 2.0), Copyright 2026 Krea AI
+# and The HuggingFace Team.
+#
+# 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.
+"""Krea2OstrisEdit -- a self-contained Hugging Face community pipeline for Krea 2
+with reference-image (edit) conditioning and Ostris AI-Toolkit LoRA loading.
+
+Everything lives in this one file so it can be hosted as a hub community
+pipeline (a model repo containing just this ``pipeline.py``):
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from PIL import Image
+
+pipe = DiffusionPipeline.from_pretrained(
+    "krea/Krea-2-Turbo",
+    custom_pipeline="ostris/Krea2OstrisEdit",
+    torch_dtype=torch.bfloat16,
+)
+pipe.to("cuda")  # or pipe.enable_model_cpu_offload() on GPUs with < ~40 GB VRAM
+
+# Load an AI-Toolkit (or already-diffusers-format) Krea 2 LoRA, e.g. the style
+# reference LoRA (generates the prompt in the style of the reference images).
+pipe.load_lora_weights(
+    "ostris/krea2_turbo_style_reference", weight_name="krea2_style_reference.safetensors"
+)
+
+image = pipe(
+    "a white yeti with horns reading a book",
+    image=Image.open("style_reference.png"),  # one reference image or a list of them
+    num_inference_steps=8,                    # Turbo defaults; the base model wants 28 / 4.5
+    guidance_scale=0.0,
+    # kv_cache=True,  # reference K/V computed once and reused every step; only for
+    #                 # LoRAs trained with AI-Toolkit's kv_cache model kwarg
+).images[0]
+image.save("output.png")
+```
+
+Reference images condition the model in two places, matching how the edit LoRAs
+are trained with Ostris AI-Toolkit (and the ComfyUI-Krea2-Ostris-Edit nodes):
+
+1. through the Qwen3-VL text encoder: each image is embedded in the user message
+   ahead of the prompt via ``Picture N: <|vision_start|><|image_pad|><|vision_end|>``
+   placeholders, so the text embeddings "see" the references;
+2. as clean VAE latents appended after the noisy image tokens in the transformer
+   sequence. They keep the flow time ``t=0`` (they are never noised) and sit on
+   rotary-position frame axis ``i + 1`` -- the Kontext-style "index" placement.
+
+Without ``image`` the pipeline is a plain Krea 2 text-to-image sampler.
+"""
+
+import math
+import os
+import re
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import PIL.Image
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import diffusers
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.loaders import PeftAdapterMixin
+from diffusers.models import AutoencoderKLQwenImage
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import randn_tensor
+
+try:
+    from transformers import AutoTokenizer, Qwen3VLModel
+except ImportError as e:  # pragma: no cover
+    raise ImportError(
+        "Krea2OstrisEdit requires a transformers version that ships Qwen3-VL "
+        "(`transformers>=4.57`). Please upgrade transformers."
+    ) from e
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+# torch>=2.5 supports grouped-query attention natively in SDPA; older versions
+# need the key/value heads repeated to the query head count.
+_SDPA_HAS_GQA = tuple(int(re.sub(r"\D.*", "", v) or 0) for v in torch.__version__.split(".")[:2]) >= (2, 5)
+
+
+# ---------------------------------------------------------------------------
+# Transformer (Krea 2 single-stream MMDiT)
+#
+# Module tree and state-dict keys match the `Krea2Transformer2DModel` checkpoint
+# layout in the `transformer/` folder of the Krea 2 hub repos, so the sharded
+# weights load directly. The forward pass additionally supports clean reference
+# tokens appended after the image tokens (`ref_seq_len`), which are modulated at
+# flow time t=0 while the text + noisy image tokens keep the real timestep.
+# ---------------------------------------------------------------------------
+
+
+class Krea2RMSNorm(nn.Module):
+    """RMSNorm with a zero-centered scale: the effective multiplier is ``1 + weight``,
+    matching the Krea 2 checkpoint format. Normalization runs in float32."""
+
+    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.float() + 1.0, eps=self.eps
+        )
+        return hidden_states.to(dtype)
+
+
+def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """Rotate interleaved (even, odd) channel pairs. ``x`` is (B, H, S, D); ``cos``/``sin``
+    are (S, D) in the repeat-interleaved layout produced by ``Krea2RotaryPosEmbed``."""
+    x_f = x.float()
+    x_rot = torch.stack((-x_f[..., 1::2], x_f[..., 0::2]), dim=-1).flatten(-2)
+    return (x_f * cos + x_rot * sin).to(x.dtype)
+
+
+class Krea2RotaryPosEmbed(nn.Module):
+    def __init__(self, theta: float, axes_dim: List[int]) -> None:
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+
+    def forward(self, ids: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        # ids: (seq_len, 3) rotary coordinates. Frequencies are computed in float64
+        # (float32 on backends without float64 support, e.g. MPS).
+        dtype = torch.float32 if ids.device.type == "mps" else torch.float64
+        angles = []
+        for i, dim in enumerate(self.axes_dim):
+            pos = ids[:, i].to(dtype)
+            freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=dtype, device=ids.device) / dim))
+            angles.append(pos[:, None] * freqs[None, :])
+        angles = torch.cat(angles, dim=-1)
+        cos = angles.cos().repeat_interleave(2, dim=-1).float()
+        sin = angles.sin().repeat_interleave(2, dim=-1).float()
+        return cos, sin
+
+
+class Krea2Attention(nn.Module):
+    """Self-attention with grouped-query projections, q/k RMSNorm, rotary embeddings
+    and a sigmoid output gate."""
+
+    def __init__(self, hidden_size: int, num_heads: int, num_kv_heads: Optional[int] = None, eps: float = 1e-5):
+        super().__init__()
+        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.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)])
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        kv_capture: Optional[list] = None,
+        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        query = self.to_q(hidden_states).unflatten(-1, (self.num_heads, self.head_dim)).transpose(1, 2)
+        key = self.to_k(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim)).transpose(1, 2)
+        value = self.to_v(hidden_states).unflatten(-1, (self.num_kv_heads, self.head_dim)).transpose(1, 2)
+        gate = self.to_gate(hidden_states)
+
+        query = self.norm_q(query)
+        key = self.norm_k(key)
+
+        if image_rotary_emb is not None:
+            cos, sin = image_rotary_emb
+            query = _apply_rotary_emb(query, cos, sin)
+            key = _apply_rotary_emb(key, cos, sin)
+
+        if kv_capture is not None:
+            kv_capture.append((key, value))
+        if kv_cache is not None:
+            # Cached reference K/V, already rotary-embedded at their original positions.
+            key = torch.cat([key, kv_cache[0].to(key.dtype)], dim=2)
+            value = torch.cat([value, kv_cache[1].to(value.dtype)], dim=2)
+
+        is_gqa = self.num_heads != self.num_kv_heads
+        if is_gqa and not _SDPA_HAS_GQA:
+            key = key.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
+            value = value.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1)
+        sdpa_kwargs = {"enable_gqa": True} if (is_gqa and _SDPA_HAS_GQA) else {}
+        hidden_states = F.scaled_dot_product_attention(query, key, value, attn_mask=attention_mask, **sdpa_kwargs)
+
+        hidden_states = hidden_states.transpose(1, 2).flatten(2)
+        hidden_states = hidden_states * torch.sigmoid(gate)
+        return self.to_out[0](hidden_states)
+
+
+class Krea2SwiGLU(nn.Module):
+    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: Optional[torch.Tensor] = 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 one text sequence:
+    ``layerwise_blocks`` attend across the layer axis per token, a linear ``projector``
+    collapses that axis, and ``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: Optional[torch.Tensor] = None):
+        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: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, int]],
+        image_rotary_emb: Tuple[torch.Tensor, torch.Tensor],
+        attention_mask: Optional[torch.Tensor] = None,
+        kv_capture: Optional[list] = None,
+        kv_cache: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        # ``temb`` is the (B, 1, 6 * hidden_size) modulation input, or a tuple
+        # ``(temb, ref_temb, split)`` for reference-image conditioning: tokens
+        # ``[:split]`` (text + noisy image) are modulated with the real timestep
+        # while tokens ``[split:]`` (clean reference tokens) use the t=0 embedding.
+        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,
+                kv_capture=kv_capture,
+                kv_cache=kv_cache,
+            )
+            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,
+            kv_capture=kv_capture,
+            kv_cache=kv_cache,
+        )
+        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 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."""
+
+    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)
+
+
+class Krea2Transformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
+    r"""
+    The Krea 2 single-stream MMDiT flow-matching backbone, extended with support for
+    clean reference-image tokens ("edit" conditioning).
+
+    Text conditioning enters as a stack of hidden states tapped from several layers of
+    the Qwen3-VL 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 (and, optionally, packed reference latents) into a single
+    ``[text, image, refs]`` sequence processed by the transformer blocks.
+
+    When ``ref_seq_len > 0``, the last ``ref_seq_len`` tokens of ``hidden_states`` are
+    clean reference tokens: they are modulated with the t=0 timestep embedding
+    (Kontext-style "index_timestep_zero") and excluded from the returned velocity.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["Krea2TransformerBlock", "Krea2TextFusionBlock", "Krea2FinalLayer"]
+    _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 precompute_ref_kv(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.Tensor,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> List[Tuple[torch.Tensor, torch.Tensor]]:
+        r"""
+        Run only the packed clean reference tokens through the transformer blocks at t=0
+        and return each block's rotary-embedded key/value pair.
+
+        Only valid for adapters trained with AI-Toolkit's ``kv_cache`` model kwarg,
+        where reference tokens attend solely to each other: their per-block K/V are
+        then independent of the timestep and of the rest of the sequence, so this
+        single pass serves every denoising step. Pass the result to
+        ``forward(..., ref_kv_cache=...)`` with the reference tokens dropped from
+        ``hidden_states`` and ``position_ids``.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, ref_seq_len, in_channels)`):
+                Packed reference latents (see the pipeline's `_pack_reference_latents`).
+            position_ids (`torch.Tensor` of shape `(ref_seq_len, 3)`):
+                The reference tokens' rotary coordinates.
+            attention_kwargs (`dict`, *optional*):
+                When it contains a `scale` entry, sets the LoRA scale, matching `forward`.
+
+        Returns:
+            A list with one `(key, value)` tuple per transformer block, each of shape
+            `(batch_size, num_key_value_heads, ref_seq_len, attention_head_dim)`.
+        """
+        lora_scale = 1.0
+        if attention_kwargs is not None:
+            lora_scale = attention_kwargs.get("scale", 1.0)
+        if USE_PEFT_BACKEND and lora_scale != 1.0:
+            scale_lora_layers(self, lora_scale)
+
+        # Clean reference tokens are always conditioned at flow time t=0.
+        timestep = torch.zeros(hidden_states.shape[0], device=hidden_states.device)
+        temb = self.time_embed(timestep, dtype=hidden_states.dtype)
+        temb_mod = self.time_mod_proj(F.gelu(temb, approximate="tanh"))
+
+        hidden_states = self.img_in(hidden_states)
+        image_rotary_emb = self.rotary_emb(position_ids)
+
+        ref_kv = []
+        for block in self.transformer_blocks:
+            captured = []
+            hidden_states = block(hidden_states, temb_mod, image_rotary_emb, kv_capture=captured)
+            ref_kv.append(captured[0])
+
+        if USE_PEFT_BACKEND and lora_scale != 1.0:
+            unscale_lora_layers(self, lora_scale)
+        return ref_kv
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        position_ids: torch.Tensor,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+        ref_seq_len: int = 0,
+        ref_kv_cache: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[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.
+            ref_seq_len (`int`, defaults to 0):
+                Number of trailing reference tokens in `hidden_states`. They receive the
+                t=0 modulation and are excluded from the output.
+            ref_kv_cache (`list[tuple[torch.Tensor, torch.Tensor]]`, *optional*):
+                Per-block reference K/V from [`~Krea2Transformer2DModel.precompute_ref_kv`].
+                When given, `hidden_states` / `position_ids` must not contain the
+                reference tokens (`ref_seq_len == 0`); each block's attention appends the
+                cached K/V as extra keys instead. Only valid for adapters trained with
+                AI-Toolkit's ``kv_cache`` model kwarg (isolated reference attention).
+            attention_kwargs (`dict`, *optional*):
+                When it contains a `scale` entry, sets the LoRA scale applied to this
+                transformer's adapters for the duration of the forward pass.
+
+        Returns:
+            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)}.")
+        if ref_kv_cache is not None and ref_seq_len > 0:
+            raise ValueError(
+                "`ref_kv_cache` replaces the reference tokens; do not also append them to "
+                "`hidden_states` (`ref_seq_len` must be 0)."
+            )
+
+        lora_scale = 1.0
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        if USE_PEFT_BACKEND and lora_scale != 1.0:
+            scale_lora_layers(self, lora_scale)
+
+        batch_size, image_seq_len, _ = hidden_states.shape  # includes ref 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"))
+
+        block_temb = temb_mod
+        if ref_seq_len > 0:
+            # Clean reference tokens are conditioned at t=0; everything else keeps t.
+            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)
+
+        # An all-True mask (no padded text tokens, e.g. any batch-of-1 call) is
+        # equivalent to no mask; passing None keeps SDPA on its fast, low-memory
+        # (flash) path instead of a mask-materializing fallback.
+        if encoder_attention_mask is not None and bool(encoder_attention_mask.all()):
+            encoder_attention_mask = None
+
+        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, :]
+            if ref_kv_cache is not None:
+                # Cached reference K/V are appended as extra (always-valid) keys.
+                ref_mask = attention_mask.new_ones((batch_size, 1, 1, ref_kv_cache[0][0].shape[2]))
+                attention_mask = torch.cat([attention_mask, ref_mask], dim=-1)
+
+        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 i, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                ckpt_func = getattr(self, "_gradient_checkpointing_func", None)
+                if ckpt_func is None:
+                    hidden_states = torch.utils.checkpoint.checkpoint(
+                        block, hidden_states, block_temb, image_rotary_emb, attention_mask, use_reentrant=False
+                    )
+                else:
+                    hidden_states = ckpt_func(block, hidden_states, block_temb, image_rotary_emb, attention_mask)
+            else:
+                hidden_states = block(
+                    hidden_states,
+                    block_temb,
+                    image_rotary_emb,
+                    attention_mask,
+                    kv_cache=ref_kv_cache[i] if ref_kv_cache is not None else None,
+                )
+
+        hidden_states = hidden_states[:, text_seq_len : text_seq_len + image_seq_len - ref_seq_len]
+        output = self.final_layer(hidden_states, temb)
+
+        if USE_PEFT_BACKEND and lora_scale != 1.0:
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+        return Transformer2DModelOutput(sample=output)
+
+
+# The Krea 2 hub repos' `model_index.json` points the `transformer` component at
+# `["diffusers", "Krea2Transformer2DModel"]`. Registering the vendored class into the
+# diffusers namespace lets `DiffusionPipeline.from_pretrained` resolve it on diffusers
+# releases that don't ship Krea 2 yet, and guarantees the loaded transformer supports
+# the reference-image forward pass this pipeline needs (the class is a numerically
+# identical superset of the upstream one for text-to-image).
+diffusers.Krea2Transformer2DModel = Krea2Transformer2DModel
+
+
+# ---------------------------------------------------------------------------
+# LoRA key conversion (Ostris AI-Toolkit / reference-trainer -> diffusers/PEFT)
+# ---------------------------------------------------------------------------
+
+
+def _convert_non_diffusers_krea2_lora_to_diffusers(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Map original `krea-ai/krea-2` module names onto `Krea2Transformer2DModel`.
+    Handles the `diffusion_model.` prefix (AI-Toolkit saves / ComfyUI) and the
+    `base_model.model.` prefix, as well as bare module names."""
+    state_dict = {
+        (k[len("base_model.model.") :] if k.startswith("base_model.model.") else k): v for k, v in state_dict.items()
+    }
+    state_dict = {
+        (k[len("diffusion_model.") :] if k.startswith("diffusion_model.") else k): v for k, v in state_dict.items()
+    }
+
+    attn_map = {"wq": "to_q", "wk": "to_k", "wv": "to_v", "wo": "to_out.0", "gate": "to_gate"}
+    ff_map = {"gate": "ff.gate", "up": "ff.up", "down": "ff.down"}
+    # The original model stores these standalone modules under abbreviated
+    # `nn.Sequential`-style names.
+    standalone_map = {
+        "first": "img_in",
+        "last.linear": "final_layer.linear",
+        "tmlp.0": "time_embed.linear_1",
+        "tmlp.2": "time_embed.linear_2",
+        "tproj.1": "time_mod_proj",
+        "txtmlp.1": "txt_in.linear_1",
+        "txtmlp.3": "txt_in.linear_2",
+        "txtfusion.projector": "text_fusion.projector",
+    }
+
+    def convert_module(module):
+        m = re.match(r"blocks\.(\d+)\.(attn|mlp)\.(\w+)$", module)
+        if m:
+            idx, kind, sub = m.groups()
+            if kind == "attn" and sub in attn_map:
+                return f"transformer_blocks.{idx}.attn.{attn_map[sub]}"
+            if kind == "mlp" and sub in ff_map:
+                return f"transformer_blocks.{idx}.{ff_map[sub]}"
+            return None
+        m = re.match(r"txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.(attn|mlp)\.(\w+)$", module)
+        if m:
+            block, idx, kind, sub = m.groups()
+            if kind == "attn" and sub in attn_map:
+                return f"text_fusion.{block}.{idx}.attn.{attn_map[sub]}"
+            if kind == "mlp" and sub in ff_map:
+                return f"text_fusion.{block}.{idx}.{ff_map[sub]}"
+            return None
+        return standalone_map.get(module)
+
+    converted_state_dict = {}
+    for key in list(state_dict):
+        match = re.search(r"\.(?:lora_[AB])\.weight$", key)
+        if match is None:
+            continue
+        diffusers_module = convert_module(key[: match.start()])
+        if diffusers_module is None:
+            continue
+        converted_state_dict[f"transformer.{diffusers_module}{key[match.start() :]}"] = state_dict.pop(key)
+
+    if len(state_dict) > 0:
+        raise ValueError(f"Could not convert LoRA keys: {sorted(state_dict.keys())}")
+
+    return converted_state_dict
+
+
+def _normalize_lora_state_dict(state_dict: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]:
+    """Normalize a Krea 2 LoRA state dict to PEFT `lora_A`/`lora_B` naming and fold any
+    `.alpha` tensors into `lora_B` so the effective scale is preserved."""
+    state_dict = {
+        k.replace(".lora_down.weight", ".lora_A.weight").replace(".lora_up.weight", ".lora_B.weight"): v
+        for k, v in state_dict.items()
+    }
+    # PEFT assumes lora_alpha == rank (scale 1.0) when no alpha is given; fold any
+    # explicit alpha into lora_B instead of plumbing network_alphas through.
+    for alpha_key in [k for k in state_dict if k.endswith(".alpha")]:
+        base = alpha_key[: -len(".alpha")]
+        a_key, b_key = base + ".lora_A.weight", base + ".lora_B.weight"
+        alpha = float(state_dict.pop(alpha_key))
+        if a_key in state_dict and b_key in state_dict:
+            rank = state_dict[a_key].shape[0]
+            if alpha != rank:
+                state_dict[b_key] = state_dict[b_key] * (alpha / rank)
+    return state_dict
+
+
+# ---------------------------------------------------------------------------
+# Pipeline
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class Krea2PipelineOutput(BaseOutput):
+    """Output class for the Krea 2 pipeline.
+
+    Args:
+        images (`list[PIL.Image.Image]` or `np.ndarray`):
+            List of `num_batches * num_images_per_prompt` denoised PIL images or a
+            numpy array of shape `(batch_size, height, width, num_channels)`.
+    """
+
+    images: Union[List[PIL.Image.Image], np.ndarray]
+
+
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 6400,
+    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
+
+
+class Krea2OstrisEditPipeline(DiffusionPipeline):
+    r"""
+    Krea 2 text-to-image / reference-image-edit pipeline with Ostris AI-Toolkit LoRA
+    loading. See the module docstring for usage.
+
+    Args:
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            Euler flow-matching scheduler configured with the Krea 2 resolution-aware
+            exponential time shift.
+        vae ([`AutoencoderKLQwenImage`]):
+            The Qwen-Image VAE (f8, 16 latent channels).
+        text_encoder ([`~transformers.Qwen3VLModel`]):
+            Qwen3-VL, including its vision tower (used to embed reference images into
+            the prompt conditioning).
+        tokenizer ([`~transformers.AutoTokenizer`]):
+            The tokenizer paired with the text encoder.
+        transformer ([`Krea2Transformer2DModel`]):
+            The Krea 2 single-stream MMDiT.
+        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.
+        is_distilled (`bool`, *optional*, defaults to `False`):
+            Whether the transformer is the few-step distilled (Turbo) checkpoint. When
+            `True`, a fixed timestep shift `mu=1.15` is used and the call defaults
+            change to `num_inference_steps=8, guidance_scale=0.0`.
+        patch_size (`int`, *optional*, defaults to 2):
+            Side length of the square patches the latents are packed into.
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+
+    # Default hub repo used to lazily build the Qwen3-VL processor that turns
+    # reference images into vision tokens (the Krea 2 repos ship only a tokenizer).
+    vl_processor_id = "Qwen/Qwen3-VL-4B-Instruct"
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLQwenImage,
+        text_encoder: Qwen3VLModel,
+        tokenizer: AutoTokenizer,
+        transformer: Krea2Transformer2DModel,
+        text_encoder_select_layers: Optional[Union[Tuple[int, ...], List[int]]] = None,
+        is_distilled: bool = False,
+        patch_size: int = 2,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            scheduler=scheduler,
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            transformer=transformer,
+        )
+        if text_encoder_select_layers is None:
+            text_encoder_select_layers = (2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35)
+        self.register_to_config(text_encoder_select_layers=tuple(text_encoder_select_layers))
+        self.text_encoder_select_layers = tuple(text_encoder_select_layers)
+        self.register_to_config(is_distilled=is_distilled)
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.register_to_config(patch_size=patch_size)
+        self.patch_size = patch_size
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * self.patch_size)
+
+        # Fixed instruction template wrapped around every prompt. The system prefix is
+        # fed through the encoder as context but its hidden states are sliced off.
+        self.prompt_template_encode_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"
+        )
+        self.prompt_template_encode_suffix = "<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 34
+
+        self._vl_processor = None
+
+    # ------------------------------------------------------------------
+    # Prompt encoding (Qwen3-VL; reference images embedded via vision tokens)
+    # ------------------------------------------------------------------
+    @property
+    def vl_processor(self):
+        """Qwen3-VL AutoProcessor, loaded lazily (only needed when reference images are
+        encoded into the prompt)."""
+        if self._vl_processor is None:
+            from transformers import AutoProcessor
+
+            self._vl_processor = AutoProcessor.from_pretrained(self.vl_processor_id)
+        return self._vl_processor
+
+    @staticmethod
+    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(self, 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 ref 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)
+                    .clamp(0, 1)
+                )
+            prepped.append(img.float())
+        return prepped
+
+    def _encode_single_prompt(
+        self,
+        prompt: str,
+        images: Optional[List[torch.Tensor]] = None,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+    ) -> torch.Tensor:
+        """Encode one prompt (optionally with reference images embedded as vision
+        tokens) into stacked Qwen3-VL hidden states of shape `(seq_len, num_text_layers,
+        text_hidden_dim)` at natural (unpadded) length."""
+        device = device or self._execution_device
+        prefix_idx = self.prompt_template_encode_start_idx
+
+        # The suffix is tokenized separately so it lands after the prompt tokens.
+        suffix_inputs = self.tokenizer([self.prompt_template_encode_suffix], return_tensors="pt").to(device)
+        suffix_ids = suffix_inputs["input_ids"]
+        suffix_mask = suffix_inputs["attention_mask"].bool()
+
+        extra_inputs = {}
+        if images:
+            # Reference images ride in the user message ahead of the prompt via named
+            # vision placeholders; the processor expands each <|image_pad|> to the
+            # image's token grid.
+            image_prompt = "".join(
+                f"Picture {i + 1}: <|vision_start|><|image_pad|><|vision_end|>" for i in range(len(images))
+            )
+            text = self.prompt_template_encode_prefix + image_prompt + prompt
+            # No truncation here: the expanded image-pad runs must stay intact.
+            inputs = self.vl_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(self.text_encoder.dtype)
+                extra_inputs[k] = v
+        else:
+            text = self.prompt_template_encode_prefix + prompt
+            inputs = self.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 = self.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 self.text_encoder_select_layers], dim=2)
+        # Drop the system-prefix tokens; what remains is (image +) prompt + suffix.
+        return hidden_states[0, prefix_idx:]
+
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        images: Optional[List[torch.Tensor]] = None,
+        num_images_per_prompt: int = 1,
+        max_sequence_length: int = 512,
+        device: Optional[torch.device] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Encode prompts (all sharing the same reference images, if any) and right-pad
+        them into a batch. Returns `(prompt_embeds, prompt_embeds_mask)` of shapes
+        `(B, L, num_text_layers, D)` and `(B, L)` (bool)."""
+        device = device or self._execution_device
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+
+        features = [self._encode_single_prompt(p, images, max_sequence_length, device) for p in prompt]
+        max_len = max(f.shape[0] for f in features)
+        embeds = features[0].new_zeros(len(features), max_len, *features[0].shape[1:])
+        mask = torch.zeros(len(features), max_len, dtype=torch.bool, device=device)
+        for i, f in enumerate(features):
+            embeds[i, : f.shape[0]] = f
+            mask[i, : f.shape[0]] = True
+
+        embeds = embeds.repeat_interleave(num_images_per_prompt, dim=0)
+        mask = mask.repeat_interleave(num_images_per_prompt, dim=0)
+        return embeds, mask
+
+    # ------------------------------------------------------------------
+    # Latent packing helpers
+    # ------------------------------------------------------------------
+    def _pack_latents(self, latents: torch.Tensor) -> torch.Tensor:
+        """(B, C, H, W) latents -> (B, H/p * W/p, C * p * p) tokens."""
+        b, c, h, w = latents.shape
+        p = self.patch_size
+        latents = latents.view(b, c, h // p, p, w // p, p)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        return latents.reshape(b, (h // p) * (w // p), c * p * p)
+
+    def _unpack_latents(self, latents: torch.Tensor, height: int, width: int) -> torch.Tensor:
+        """(B, L, C * p * p) tokens -> (B, C, 1, H, W) latents (frame dim for the VAE)."""
+        batch_size, _, channels = latents.shape
+        p = self.patch_size
+        h = p * (int(height) // (self.vae_scale_factor * p))
+        w = p * (int(width) // (self.vae_scale_factor * p))
+        latents = latents.view(batch_size, h // p, w // p, channels // (p * p), p, p)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+        return latents.reshape(batch_size, channels // (p * p), 1, h, w)
+
+    def prepare_latents(self, batch_size, num_channels_latents, height, width, dtype, device, generator, latents=None):
+        if latents is not None:
+            return latents.to(device=device, dtype=dtype)
+        shape = (batch_size, num_channels_latents, height // self.vae_scale_factor, width // self.vae_scale_factor)
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+        return self._pack_latents(latents)
+
+    def _encode_reference_latents(
+        self,
+        images: List[torch.Tensor],
+        max_pixels: int,
+        generator: Optional[torch.Generator],
+        device: torch.device,
+    ) -> 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."""
+        snap = self.vae_scale_factor * self.patch_size
+        vae_dtype = self.vae.dtype
+
+        latents_mean = torch.tensor(self.vae.config.latents_mean).view(1, self.vae.config.z_dim, 1, 1, 1)
+        latents_std = torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 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 = self.vae.encode(img).latent_dist.sample(generator)
+            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(
+        self,
+        ref_latents: List[torch.Tensor],
+        device: torch.device,
+        dtype: torch.dtype,
+        placements: Optional[List[Dict[str, Any]]] = None,
+        target_grid_size: Optional[Tuple[int, int]] = None,
+    ) -> 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 (Kontext "index" placement)."""
+        p = self.patch_size
+        tokens, position_ids = [], []
+        for i, ref in enumerate(ref_latents):
+            ref = ref.unsqueeze(0).to(device, dtype)
+            tokens.append(self._pack_latents(ref))
+            _, _, h, w = ref.shape
+            ids = torch.zeros(h // p, w // p, 3, device=device)
+            ids[..., 0] = i + 1
+            placement = placements[i] if placements is not None else None
+            if placement is None:
+                ids[..., 1] = torch.arange(h // p, device=device)[:, None]
+                ids[..., 2] = torch.arange(w // p, device=device)[None, :]
+            else:
+                if target_grid_size is None:
+                    raise ValueError("`target_grid_size` is required for registered references.")
+                bbox = placement.get("bbox_normalized")
+                if bbox is None or len(bbox) != 4:
+                    raise ValueError("A registered reference requires bbox_normalized=[x0,y0,x1,y1].")
+                x0, y0, x1, y1 = (float(value) for value in bbox)
+                if not (0 <= x0 < x1 <= 1 and 0 <= y0 < y1 <= 1):
+                    raise ValueError(f"Invalid normalized reference bbox: {bbox}")
+                target_h, target_w = target_grid_size
+                ref_h, ref_w = h // p, w // p
+                ys = y0 * target_h + (torch.arange(ref_h, device=device) + 0.5) * (
+                    (y1 - y0) * target_h / ref_h
+                ) - 0.5
+                xs = x0 * target_w + (torch.arange(ref_w, device=device) + 0.5) * (
+                    (x1 - x0) * target_w / ref_w
+                ) - 0.5
+                ids[..., 1] = ys[:, None]
+                ids[..., 2] = xs[None, :]
+            position_ids.append(ids.reshape(-1, 3))
+        return torch.cat(tokens, dim=1), torch.cat(position_ids, dim=0)
+
+    @staticmethod
+    def prepare_position_ids(text_seq_len: int, grid_height: int, grid_width: int, device: torch.device):
+        """Rotary coordinates for the `[text, image]` sequence: text tokens sit at the
+        origin, image tokens carry their `(0, h, w)` latent-grid coordinates."""
+        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)
+        return torch.cat([text_ids, image_ids], dim=0)
+
+    # ------------------------------------------------------------------
+    # LoRA loading (Ostris AI-Toolkit / ComfyUI / diffusers formats)
+    # ------------------------------------------------------------------
+    def load_lora_weights(
+        self,
+        pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+        weight_name: Optional[str] = None,
+        adapter_name: str = "default",
+        **kwargs,
+    ):
+        r"""
+        Load a Krea 2 LoRA into the transformer.
+
+        Accepts a state dict, a local `.safetensors` file or directory, or a hub repo id
+        (with `weight_name` selecting the file when the repo holds several). Handles
+        Ostris AI-Toolkit / ComfyUI key layouts (`diffusion_model.blocks...` with
+        `lora_A`/`lora_B` or `lora_down`/`lora_up`) as well as already-converted
+        diffusers-format state dicts (`transformer.transformer_blocks...`).
+        """
+        if isinstance(pretrained_model_name_or_path_or_dict, dict):
+            state_dict = dict(pretrained_model_name_or_path_or_dict)
+        else:
+            from safetensors.torch import load_file
+
+            path = str(pretrained_model_name_or_path_or_dict)
+            if os.path.isfile(path):
+                file_path = path
+            elif os.path.isdir(path):
+                if weight_name is None:
+                    candidates = [f for f in os.listdir(path) if f.endswith(".safetensors")]
+                    if len(candidates) != 1:
+                        raise ValueError(
+                            f"Could not pick a LoRA file in {path}: found {candidates}. Pass `weight_name`."
+                        )
+                    weight_name = candidates[0]
+                file_path = os.path.join(path, weight_name)
+            else:
+                from huggingface_hub import hf_hub_download, list_repo_files
+
+                if weight_name is None:
+                    candidates = [
+                        f for f in list_repo_files(path, token=kwargs.get("token", None)) if f.endswith(".safetensors")
+                    ]
+                    if len(candidates) != 1:
+                        raise ValueError(
+                            f"Could not pick a LoRA file in hub repo {path}: found {candidates}. Pass `weight_name`."
+                        )
+                    weight_name = candidates[0]
+                file_path = hf_hub_download(path, weight_name, token=kwargs.get("token", None))
+            state_dict = load_file(file_path)
+
+        state_dict = _normalize_lora_state_dict(state_dict)
+        if not any(k.startswith("transformer.") for k in state_dict):
+            state_dict = _convert_non_diffusers_krea2_lora_to_diffusers(state_dict)
+
+        self.transformer.load_lora_adapter(state_dict, prefix="transformer", adapter_name=adapter_name)
+
+    def unload_lora_weights(self):
+        """Remove all loaded LoRA adapters from the transformer."""
+        transformer = self.transformer
+        if hasattr(transformer, "unload_lora"):
+            transformer.unload_lora()
+        elif getattr(transformer, "peft_config", None):
+            transformer.delete_adapters(list(transformer.peft_config.keys()))
+
+    def fuse_lora(self, lora_scale: float = 1.0, adapter_names: Optional[List[str]] = None, **kwargs):
+        """Fuse the loaded LoRA weights into the transformer for adapter-free inference."""
+        self.transformer.fuse_lora(lora_scale=lora_scale, adapter_names=adapter_names, **kwargs)
+
+    def unfuse_lora(self, **kwargs):
+        self.transformer.unfuse_lora(**kwargs)
+
+    def set_adapters(self, adapter_names: Union[str, List[str]], weights: Optional[Union[float, List[float]]] = None):
+        """Activate (and optionally weight) specific loaded LoRA adapters."""
+        self.transformer.set_adapters(adapter_names, weights)
+
+    # ------------------------------------------------------------------
+    # Generation
+    # ------------------------------------------------------------------
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def do_classifier_free_guidance(self):
+        return self._guidance_scale > 0
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str], None] = None,
+        image: Union[PIL.Image.Image, np.ndarray, torch.Tensor, List, None] = None,
+        negative_prompt: Union[str, List[str], None] = None,
+        height: int = 1024,
+        width: int = 1024,
+        num_inference_steps: Optional[int] = None,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: Optional[float] = None,
+        num_images_per_prompt: int = 1,
+        generator: Union[torch.Generator, List[torch.Generator], None] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+        reference_max_pixels: int = 1024 * 1024,
+        reference_placements: Optional[List[Dict[str, Any]]] = None,
+        vl_image_max_pixels: int = 384 * 384,
+        encode_reference_in_prompt: bool = True,
+        kv_cache: bool = False,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        max_sequence_length: int = 512,
+    ):
+        r"""
+        Generate images from a prompt, optionally conditioned on reference images.
+
+        Args:
+            prompt (`str` or `list[str]`):
+                The prompt(s) to guide generation. For edits, describe the change (e.g.
+                "make the sky purple").
+            image (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor` or a list of them, *optional*):
+                Reference image(s). They are encoded into the prompt conditioning via
+                the Qwen3-VL vision tower and appended to the transformer sequence as
+                clean VAE latents at t=0. References keep their own aspect ratio; the
+                output size is set by `height`/`width` independently.
+            negative_prompt (`str` or `list[str]`, *optional*):
+                Prompt(s) not to guide generation; ignored when `guidance_scale <= 0`.
+            height / width (`int`, defaults to 1024):
+                Output size in pixels; rounded up to a multiple of 16 if needed.
+            num_inference_steps (`int`, *optional*):
+                Denoising steps. Defaults to 8 for a distilled (Turbo) checkpoint and 28
+                otherwise.
+            sigmas (`list[float]`, *optional*):
+                Custom sigma grid for the scheduler.
+            guidance_scale (`float`, *optional*):
+                Krea 2 CFG convention: velocity is `cond + scale * (cond - uncond)` and
+                guidance is enabled whenever `scale > 0` (equals standard CFG with scale
+                `1 + scale`). Defaults to 0.0 for a distilled checkpoint and 4.5
+                otherwise.
+            num_images_per_prompt (`int`, defaults to 1):
+                Number of images per prompt.
+            generator (`torch.Generator` or `list[torch.Generator]`, *optional*):
+                RNG for deterministic generation.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated packed noisy latents `(B, image_seq_len, in_channels)`.
+            prompt_embeds / prompt_embeds_mask (`torch.Tensor`, *optional*):
+                Pre-computed text conditioning `(B, L, num_text_layers, D)` and its
+                bool mask `(B, L)`; skips prompt encoding when given.
+            negative_prompt_embeds / negative_prompt_embeds_mask (`torch.Tensor`, *optional*):
+                Same, for the negative prompt.
+            reference_max_pixels (`int`, defaults to `1024 * 1024`):
+                Pixel budget each reference image is downscaled to fit before VAE
+                encoding (never upscaled).
+            vl_image_max_pixels (`int`, defaults to `384 * 384`):
+                Pixel budget for the (coarse) Qwen3-VL view of each reference image.
+            encode_reference_in_prompt (`bool`, defaults to `True`):
+                Whether reference images are also embedded into the text conditioning
+                through the Qwen3-VL vision tower (matches AI-Toolkit edit training).
+            kv_cache (`bool`, defaults to `False`):
+                Cache the reference tokens' attention K/V: they are precomputed
+                in a single t=0 pass and reused on every denoising step, so the
+                reference tokens never ride along in the per-step sequence --
+                faster, especially with CFG or many steps. The LoRA must be
+                trained with AI-Toolkit's ``kv_cache`` model kwarg (reference
+                tokens attend only to each other) for this to work properly;
+                leave off for normally trained edit LoRAs.
+            output_type (`str`, defaults to `"pil"`):
+                `"pil"`, `"np"`, `"pt"` or `"latent"`.
+            return_dict (`bool`, defaults to `True`):
+                Whether to return a [`Krea2PipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                Forwarded to the transformer; a `scale` entry sets the LoRA scale.
+            max_sequence_length (`int`, defaults to 512):
+                Maximum prompt token length (truncation only; no fixed padding).
+
+        Returns:
+            [`Krea2PipelineOutput`] or `tuple`: the generated images.
+        """
+        if num_inference_steps is None:
+            num_inference_steps = 8 if self.config.is_distilled else 28
+        if guidance_scale is None:
+            guidance_scale = 0.0 if self.config.is_distilled else 4.5
+
+        multiple = self.vae_scale_factor * self.patch_size
+        if height % multiple != 0 or width % multiple != 0:
+            rounded_height = ((height + multiple - 1) // multiple) * multiple
+            rounded_width = ((width + multiple - 1) // multiple) * multiple
+            logger.warning(
+                f"`height` and `width` must be multiples of {multiple}; rounding up from {height}x{width} to"
+                f" {rounded_height}x{rounded_width}."
+            )
+            height, width = rounded_height, rounded_width
+
+        if prompt is None and prompt_embeds is None:
+            raise ValueError("Provide either `prompt` or `prompt_embeds`.")
+        if prompt_embeds is not None and prompt_embeds_mask is None:
+            raise ValueError("`prompt_embeds` requires `prompt_embeds_mask`.")
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            raise ValueError("`negative_prompt_embeds` requires `negative_prompt_embeds_mask`.")
+
+        self._guidance_scale = guidance_scale
+
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None:
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        transformer_dtype = self.transformer.dtype
+
+        # 1. Normalize reference images to a list of [0, 1] CHW tensors.
+        ref_images = None
+        if image is not None:
+            image_list = image if isinstance(image, (list, tuple)) else [image]
+            ref_images = [self._to_chw_tensor(img) for img in image_list]
+
+        # 2. Encode the prompt(s). With references, the coarse VL view of each image is
+        # embedded in the user message so the text conditioning "sees" them.
+        vl_images = None
+        if ref_images is not None and encode_reference_in_prompt:
+            vl_images = self._prep_vl_images([img.to(device) for img in ref_images], vl_image_max_pixels)
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+                prompt, vl_images, num_images_per_prompt, max_sequence_length, device
+            )
+        prompt_embeds = prompt_embeds.to(transformer_dtype)
+
+        if self.do_classifier_free_guidance:
+            if negative_prompt_embeds is None:
+                negative_prompt = negative_prompt if negative_prompt is not None else ""
+                if isinstance(negative_prompt, str):
+                    negative_prompt = [negative_prompt] * batch_size
+                negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                    negative_prompt, vl_images, num_images_per_prompt, max_sequence_length, device
+                )
+            negative_prompt_embeds = negative_prompt_embeds.to(transformer_dtype)
+
+        # 3. Prepare the noisy latents (kept in float32 across scheduler steps).
+        num_channels_latents = self.transformer.config.in_channels // (self.patch_size**2)
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            torch.float32,
+            device,
+            generator,
+            latents,
+        )
+        grid_height = height // (self.vae_scale_factor * self.patch_size)
+        grid_width = width // (self.vae_scale_factor * self.patch_size)
+
+        # 4. Encode + pack reference latents (shared across the batch) and build the
+        # combined rotary coordinates.
+        ref_tokens, ref_seq_len = None, 0
+        neg_position_ids = None
+        position_ids = self.prepare_position_ids(prompt_embeds.shape[1], grid_height, grid_width, device)
+        if self.do_classifier_free_guidance:
+            neg_position_ids = self.prepare_position_ids(
+                negative_prompt_embeds.shape[1], grid_height, grid_width, device
+            )
+        ref_kv = None
+        if ref_images is not None:
+            if reference_placements is not None and len(reference_placements) != len(ref_images):
+                raise ValueError("`reference_placements` must match the number of reference images.")
+            ref_latents = self._encode_reference_latents(ref_images, reference_max_pixels, generator, device)
+            ref_tokens, ref_position_ids = self._pack_reference_latents(
+                ref_latents,
+                device,
+                transformer_dtype,
+                placements=reference_placements,
+                target_grid_size=(grid_height, grid_width),
+            )
+            ref_seq_len = ref_tokens.shape[1]
+            ref_tokens = ref_tokens.expand(latents.shape[0], -1, -1)
+            if kv_cache:
+                # Precompute pass: the refs alone run through the blocks once at t=0
+                # and every denoising step reuses their per-block K/V, so the ref
+                # tokens are dropped from the per-step sequence entirely.
+                ref_kv = self.transformer.precompute_ref_kv(ref_tokens, ref_position_ids, attention_kwargs)
+                ref_tokens, ref_seq_len = None, 0
+            else:
+                position_ids = torch.cat([position_ids, ref_position_ids], dim=0)
+                if neg_position_ids is not None:
+                    neg_position_ids = torch.cat([neg_position_ids, ref_position_ids], dim=0)
+
+        # 5. Prepare timesteps. The distilled (Turbo) checkpoint was trained at a fixed
+        # exponential time shift mu=1.15; the base checkpoint interpolates mu from the
+        # image token count.
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        if self.config.is_distilled:
+            mu = 1.15
+        else:
+            mu = calculate_shift(
+                grid_height * grid_width,
+                self.scheduler.config.get("base_image_seq_len", 256),
+                self.scheduler.config.get("max_image_seq_len", 6400),
+                self.scheduler.config.get("base_shift", 0.5),
+                self.scheduler.config.get("max_shift", 1.15),
+            )
+        self.scheduler.set_timesteps(sigmas=sigmas, device=device, mu=mu)
+        timesteps = self.scheduler.timesteps
+        self.scheduler.set_begin_index(0)
+
+        # 6. Denoising loop (Euler flow ODE integration via the scheduler).
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for t in timesteps:
+                timestep = (t / self.scheduler.config.num_train_timesteps).expand(latents.shape[0]).to(
+                    transformer_dtype
+                )
+
+                model_input = latents.to(transformer_dtype)
+                if ref_tokens is not None:
+                    model_input = torch.cat([model_input, ref_tokens], dim=1)
+
+                noise_pred = self.transformer(
+                    hidden_states=model_input,
+                    encoder_hidden_states=prompt_embeds,
+                    timestep=timestep,
+                    position_ids=position_ids,
+                    encoder_attention_mask=prompt_embeds_mask,
+                    ref_seq_len=ref_seq_len,
+                    ref_kv_cache=ref_kv,
+                    attention_kwargs=attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if self.do_classifier_free_guidance:
+                    neg_noise_pred = self.transformer(
+                        hidden_states=model_input,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        timestep=timestep,
+                        position_ids=neg_position_ids,
+                        encoder_attention_mask=negative_prompt_embeds_mask,
+                        ref_seq_len=ref_seq_len,
+                        ref_kv_cache=ref_kv,
+                        attention_kwargs=attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred + guidance_scale * (noise_pred - neg_noise_pred)
+
+                latents = self.scheduler.step(noise_pred.float(), t, latents, return_dict=False)[0]
+                progress_bar.update()
+
+        # 7. Decode latents.
+        if output_type == "latent":
+            image_out = latents
+        else:
+            latents = self._unpack_latents(latents, height, width).to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = (
+                torch.tensor(self.vae.config.latents_std)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents = latents * latents_std + latents_mean
+            image_out = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+            image_out = self.image_processor.postprocess(image_out, output_type=output_type)
+
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image_out,)
+        return Krea2PipelineOutput(images=image_out)