Add from_single_file() support for Ideogram4 and Krea2 transformers

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

@@ -14,6 +14,19 @@ specific language governing permissions and limitations under the License.
 
 A transformer for image-like data from [Ideogram 4](https://github.com/ideogram-oss/ideogram-4). 
 
+## Loading single-file checkpoints
+
+`Ideogram4Transformer2DModel` supports loading from single-file checkpoints via [`~loaders.FromOriginalModelMixin.from_single_file`]:
+
+```python
+from diffusers import Ideogram4Transformer2DModel
+
+transformer = Ideogram4Transformer2DModel.from_single_file(
+    "path/to/ideogram4_transformer.safetensors",
+    torch_dtype=torch.bfloat16,
+)
+```
+
 ## Ideogram4Transformer2DModel
 
 [[autodoc]] Ideogram4Transformer2DModel

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

@@ -14,6 +14,19 @@ specific language governing permissions and limitations under the License.
 
 The single-stream MMDiT flow-matching transformer used by [Krea 2](https://github.com/krea-ai/krea-2).
 
+## Loading single-file checkpoints
+
+`Krea2Transformer2DModel` supports loading from single-file checkpoints via [`~loaders.FromOriginalModelMixin.from_single_file`]:
+
+```python
+from diffusers import Krea2Transformer2DModel
+
+transformer = Krea2Transformer2DModel.from_single_file(
+    "path/to/krea2_transformer.safetensors",
+    torch_dtype=torch.bfloat16,
+)
+```
+
 ## Krea2Transformer2DModel
 
 [[autodoc]] Krea2Transformer2DModel

src/diffusers/loaders/single_file_model.py

@@ -42,6 +42,8 @@
     convert_flux_transformer_checkpoint_to_diffusers,
     convert_hidream_transformer_to_diffusers,
     convert_hunyuan_video_transformer_to_diffusers,
+    convert_ideogram4_transformer_checkpoint_to_diffusers,
+    convert_krea2_transformer_checkpoint_to_diffusers,
     convert_ldm_unet_checkpoint,
     convert_ldm_vae_checkpoint,
     convert_ltx2_audio_vae_to_diffusers,
@@ -215,6 +217,14 @@
         "checkpoint_mapping_fn": lambda checkpoint, **kwargs: checkpoint,
         "default_subfolder": "transformer",
     },
+    "Ideogram4Transformer2DModel": {
+        "checkpoint_mapping_fn": convert_ideogram4_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
+    "Krea2Transformer2DModel": {
+        "checkpoint_mapping_fn": convert_krea2_transformer_checkpoint_to_diffusers,
+        "default_subfolder": "transformer",
+    },
 }
 
 

src/diffusers/loaders/single_file_utils.py

@@ -149,6 +149,12 @@
         "net.pos_embedder.dim_spatial_range",
     ],
     "flux2": ["model.diffusion_model.single_stream_modulation.lin.weight", "single_stream_modulation.lin.weight"],
+    "ideogram4": [
+        "diffusion_model.llm_cond_proj.weight",
+        "conditional_transformer.llm_cond_proj.weight",
+        "llm_cond_proj.weight",
+    ],
+    "krea2": ["diffusion_model.txtfusion.projector.weight", "txtfusion.projector.weight"],
     "ltx2": [
         "model.diffusion_model.av_ca_a2v_gate_adaln_single.emb.timestep_embedder.linear_1.weight",
         "vae.per_channel_statistics.mean-of-means",
@@ -237,6 +243,8 @@
     "z-image-turbo-controlnet-2.0": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.0"},
     "z-image-turbo-controlnet-2.1": {"pretrained_model_name_or_path": "hlky/Z-Image-Turbo-Fun-Controlnet-Union-2.1"},
     "ltx2-dev": {"pretrained_model_name_or_path": "Lightricks/LTX-2"},
+    "ideogram4": {"pretrained_model_name_or_path": "ideogram-ai/ideogram-v4"},
+    "krea2": {"pretrained_model_name_or_path": "krea/Krea-2-Turbo"},
 }
 
 # Use to configure model sample size when original config is provided
@@ -683,6 +691,12 @@ def infer_diffusers_model_type(checkpoint):
     elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux2"]):
         model_type = "flux-2-dev"
 
+    elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["ideogram4"]):
+        model_type = "ideogram4"
+
+    elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["krea2"]):
+        model_type = "krea2"
+
     elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["flux"]):
         if any(
             g in checkpoint for g in ["guidance_in.in_layer.bias", "model.diffusion_model.guidance_in.in_layer.bias"]
@@ -4180,3 +4194,89 @@ def convert_ernie_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs
             checkpoint[k.replace("model.diffusion_model.", "")] = checkpoint.pop(k)
 
     return checkpoint
+
+
+def convert_ideogram4_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
+    converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())}
+
+    for prefix in ("diffusion_model.", "conditional_transformer."):
+        if any(k.startswith(prefix) for k in converted_state_dict):
+            converted_state_dict = {k.removeprefix(prefix): v for k, v in converted_state_dict.items()}
+            break
+
+    # attention.o -> attention.to_out.0
+    for key in list(converted_state_dict.keys()):
+        if ".attention.o." in key:
+            new_key = key.replace(".attention.o.", ".attention.to_out.0.")
+            converted_state_dict[new_key] = converted_state_dict.pop(key)
+
+    # Split fused attention.qkv -> to_q / to_k / to_v
+    for key in list(converted_state_dict.keys()):
+        if ".attention.qkv.weight" in key:
+            fused = converted_state_dict.pop(key)
+            q, k, v = torch.chunk(fused, 3, dim=0)
+            base = key[: key.index(".attention.qkv.weight")]
+            converted_state_dict[f"{base}.attention.to_q.weight"] = q.contiguous()
+            converted_state_dict[f"{base}.attention.to_k.weight"] = k.contiguous()
+            converted_state_dict[f"{base}.attention.to_v.weight"] = v.contiguous()
+
+    return converted_state_dict
+
+
+def convert_krea2_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
+    converted_state_dict = {key: checkpoint.pop(key) for key in list(checkpoint.keys())}
+
+    converted_state_dict = {k.removeprefix("diffusion_model."): v for k, v in converted_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"}
+
+    def remap_key(key):
+        # transformer blocks: blocks.{i}.*
+        m = re.match(r"^blocks\.(\d+)\.(.+)$", key)
+        if m:
+            idx, tail = m.groups()
+            if tail.startswith("attn."):
+                sub_m = re.match(r"^attn\.(\w+)(.*)$", tail)
+                if sub_m:
+                    sub, rest = sub_m.groups()
+                    if sub in ATTN_MAP:
+                        return f"transformer_blocks.{idx}.attn.{ATTN_MAP[sub]}{rest}"
+            if tail.startswith("mlp."):
+                sub_m = re.match(r"^mlp\.(\w+)(.*)$", tail)
+                if sub_m:
+                    sub, rest = sub_m.groups()
+                    if sub in FF_MAP:
+                        return f"transformer_blocks.{idx}.{FF_MAP[sub]}{rest}"
+            # norm1, norm2, scale_shift_table and any unmatched sub-keys
+            return f"transformer_blocks.{idx}.{tail}"
+
+        # text fusion blocks: txtfusion.(layerwise_blocks|refiner_blocks).{i}.(attn|mlp).*
+        m = re.match(r"^txtfusion\.(layerwise_blocks|refiner_blocks)\.(\d+)\.(.+)$", key)
+        if m:
+            block_group, idx, tail = m.groups()
+            if tail.startswith("attn."):
+                sub_m = re.match(r"^attn\.(\w+)(.*)$", tail)
+                if sub_m:
+                    sub, rest = sub_m.groups()
+                    if sub in ATTN_MAP:
+                        return f"text_fusion.{block_group}.{idx}.attn.{ATTN_MAP[sub]}{rest}"
+            if tail.startswith("mlp."):
+                sub_m = re.match(r"^mlp\.(\w+)(.*)$", tail)
+                if sub_m:
+                    sub, rest = sub_m.groups()
+                    if sub in FF_MAP:
+                        return f"text_fusion.{block_group}.{idx}.{FF_MAP[sub]}{rest}"
+            return f"text_fusion.{block_group}.{idx}.{tail}"
+
+        # top-level txtfusion.* (e.g. txtfusion.projector)
+        if key.startswith("txtfusion."):
+            return "text_fusion." + key[len("txtfusion.") :]
+
+        return key
+
+    result = {}
+    for key, value in converted_state_dict.items():
+        result[remap_key(key)] = value
+
+    return result

src/diffusers/models/transformers/transformer_krea2.py

@@ -21,7 +21,7 @@
 import torch.nn.functional as F
 
 from ...configuration_utils import ConfigMixin, register_to_config
-from ...loaders import PeftAdapterMixin
+from ...loaders import FromOriginalModelMixin, PeftAdapterMixin
 from ...utils import apply_lora_scale, logging
 from ...utils.torch_utils import maybe_adjust_dtype_for_device
 from ..attention import AttentionMixin, AttentionModuleMixin
@@ -327,7 +327,7 @@ def forward(self, ids: torch.Tensor) -> torch.Tensor:
         return freqs_cos, freqs_sin
 
 
-class Krea2Transformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin):
+class Krea2Transformer2DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin, FromOriginalModelMixin):
     r"""
     The single-stream MMDiT flow-matching backbone used by the Krea 2 pipeline.