Merge branch 'main' into ideogram4-lora-training

Author: dg845

docs/source/en/_toctree.yml

@@ -355,6 +355,8 @@
         title: Ideogram4Transformer2DModel
       - local: api/models/transformer_joyimage
         title: JoyImageEditTransformer3DModel
+      - local: api/models/krea2_transformer2d
+        title: Krea2Transformer2DModel
       - local: api/models/latte_transformer3d
         title: LatteTransformer3DModel
       - local: api/models/longcat_image_transformer2d
@@ -563,6 +565,8 @@
         title: Kandinsky 5.0 Image
       - local: api/pipelines/kolors
         title: Kolors
+      - local: api/pipelines/krea2
+        title: Krea 2
       - local: api/pipelines/latent_consistency_models
         title: Latent Consistency Models
       - local: api/pipelines/latent_diffusion

docs/source/en/api/models/krea2_transformer2d.md

@@ -0,0 +1,19 @@
+<!--Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# Krea2Transformer2DModel
+
+The single-stream MMDiT flow-matching transformer used by [Krea 2](https://github.com/krea-ai/krea-2).
+
+## Krea2Transformer2DModel
+
+[[autodoc]] Krea2Transformer2DModel

docs/source/en/api/pipelines/krea2.md

@@ -0,0 +1,61 @@
+<!--Copyright 2026 Krea AI and The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# Krea 2
+
+Krea 2 (K2) is a flow-matching text-to-image model built around a single-stream MMDiT with grouped-query attention. A
+Qwen3-VL text encoder provides the conditioning: instead of the last hidden state, hidden states from twelve decoder
+layers are tapped per token and fused inside the transformer by a small text-fusion stage. Images are decoded with the
+Qwen-Image VAE.
+
+Two checkpoints are released, sharing the same architecture but with different recommended sampler settings:
+
+- **Base (midtrain)** — use the full sampler with classifier-free guidance: `num_inference_steps=28`,
+  `guidance_scale=4.5`.
+- **TDM (distilled)** — distilled for few-step sampling, run with `num_inference_steps=8` and guidance disabled
+  (`guidance_scale=0.0`).
+
+`guidance_scale` follows the Krea 2 convention: the velocity is computed as `cond + guidance_scale * (cond - uncond)`
+and guidance is enabled whenever `guidance_scale > 0` (this equals the usual CFG formulation with scale
+`1 + guidance_scale`).
+
+## Text-to-image
+
+```python
+import torch
+from diffusers import Krea2Pipeline
+
+# Load from a local directory produced by the Krea 2 conversion (no hub repo yet).
+pipe = Krea2Pipeline.from_pretrained("path/to/krea2-diffusers", torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+prompt = "a fox in the snow"
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    num_inference_steps=28,
+    guidance_scale=4.5,
+    generator=torch.Generator("cuda").manual_seed(0),
+).images[0]
+image.save("krea2.png")
+```
+
+## Krea2Pipeline
+
+[[autodoc]] Krea2Pipeline
+  - all
+  - __call__
+
+## Krea2PipelineOutput
+
+[[autodoc]] pipelines.krea2.pipeline_output.Krea2PipelineOutput

src/diffusers/__init__.py

@@ -277,6 +277,7 @@
             "JoyImageEditTransformer3DModel",
             "Kandinsky3UNet",
             "Kandinsky5Transformer3DModel",
+            "Krea2Transformer2DModel",
             "LatteTransformer3DModel",
             "LongCatAudioDiTTransformer",
             "LongCatAudioDiTVae",
@@ -646,6 +647,7 @@
             "KandinskyV22Pipeline",
             "KandinskyV22PriorEmb2EmbPipeline",
             "KandinskyV22PriorPipeline",
+            "Krea2Pipeline",
             "LatentConsistencyModelImg2ImgPipeline",
             "LatentConsistencyModelPipeline",
             "LattePipeline",
@@ -1137,6 +1139,7 @@
             JoyImageEditTransformer3DModel,
             Kandinsky3UNet,
             Kandinsky5Transformer3DModel,
+            Krea2Transformer2DModel,
             LatteTransformer3DModel,
             LongCatAudioDiTTransformer,
             LongCatAudioDiTVae,
@@ -1481,6 +1484,7 @@
             KandinskyV22Pipeline,
             KandinskyV22PriorEmb2EmbPipeline,
             KandinskyV22PriorPipeline,
+            Krea2Pipeline,
             LatentConsistencyModelImg2ImgPipeline,
             LatentConsistencyModelPipeline,
             LattePipeline,

src/diffusers/loaders/lora_conversion_utils.py

@@ -551,11 +551,19 @@ def assign_remaining_weights(assignments, source):
                 for target_fmt, source_fmt, transform in assignments:
                     target_key = target_fmt.format(lora_key=lora_key)
                     source_key = source_fmt.format(orig_lora_key=orig_lora_key)
-                    value = source.pop(source_key)
-                    if transform:
+                    value = source.pop(source_key, None)
+                    if value is None:
+                        continue
+                    if transform and lora_key == "lora_B":
                         value = transform(value)
                     ait_sd[target_key] = value
 
+            # Consume any leftover final_layer alpha keys so they don't
+            # reach the remaining_keys guard and cause a false "Incompatible keys" error.
+            for key in list(source.keys()):
+                if "final_layer" in key and key.endswith(".alpha"):
+                    source.pop(key)
+
         if any("guidance_in" in k for k in sds_sd):
             _convert_to_ai_toolkit(
                 sds_sd,

src/diffusers/models/__init__.py

@@ -122,6 +122,7 @@
     _import_structure["transformers.transformer_ideogram4"] = ["Ideogram4Transformer2DModel"]
     _import_structure["transformers.transformer_joyimage"] = ["JoyImageEditTransformer3DModel"]
     _import_structure["transformers.transformer_kandinsky"] = ["Kandinsky5Transformer3DModel"]
+    _import_structure["transformers.transformer_krea2"] = ["Krea2Transformer2DModel"]
     _import_structure["transformers.transformer_longcat_audio_dit"] = ["LongCatAudioDiTTransformer"]
     _import_structure["transformers.transformer_longcat_image"] = ["LongCatImageTransformer2DModel"]
     _import_structure["transformers.transformer_ltx"] = ["LTXVideoTransformer3DModel"]
@@ -255,6 +256,7 @@
             Ideogram4Transformer2DModel,
             JoyImageEditTransformer3DModel,
             Kandinsky5Transformer3DModel,
+            Krea2Transformer2DModel,
             LatteTransformer3DModel,
             LongCatAudioDiTTransformer,
             LongCatImageTransformer2DModel,

src/diffusers/models/transformers/__init__.py

@@ -43,6 +43,7 @@
     from .transformer_ideogram4 import Ideogram4Transformer2DModel
     from .transformer_joyimage import JoyImageEditTransformer3DModel
     from .transformer_kandinsky import Kandinsky5Transformer3DModel
+    from .transformer_krea2 import Krea2Transformer2DModel
     from .transformer_longcat_audio_dit import LongCatAudioDiTTransformer
     from .transformer_longcat_image import LongCatImageTransformer2DModel
     from .transformer_ltx import LTXVideoTransformer3DModel