diff --git a/docs/source/en/api/models/ideogram4_transformer2d.md b/docs/source/en/api/models/ideogram4_transformer2d.md index 5cf0c1e2b0d2..f5f6cf1fc31f 100644 --- a/docs/source/en/api/models/ideogram4_transformer2d.md +++ b/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 diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 610e377d3ab9..4410dbf96cdd 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -43,6 +43,7 @@ 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, @@ -201,6 +202,10 @@ "checkpoint_mapping_fn": convert_qwen_image21_transformer_checkpoint_to_diffusers, "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", diff --git a/src/diffusers/loaders/single_file_utils.py b/src/diffusers/loaders/single_file_utils.py index e65e0be50f62..355bac29cb49 100644 --- a/src/diffusers/loaders/single_file_utils.py +++ b/src/diffusers/loaders/single_file_utils.py @@ -165,6 +165,11 @@ "audio_vae.per_channel_statistics.mean-of-means", ], "qwen-image-2.1": ["model.diffusion_model.txt_in.text_norm.weight", "txt_in.text_norm.weight"], + "ideogram4": [ + "diffusion_model.llm_cond_proj.weight", + "conditional_transformer.llm_cond_proj.weight", + "llm_cond_proj.weight", + ], "krea2": ["model.diffusion_model.txtfusion.projector.weight", "txtfusion.projector.weight"], } @@ -253,6 +258,7 @@ "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"}, "qwen-image-2.1": {"pretrained_model_name_or_path": "Qwen/Qwen-Image-2.1"}, + "ideogram4": {"pretrained_model_name_or_path": "ideogram-ai/ideogram-v4"}, "krea2": {"pretrained_model_name_or_path": "krea/Krea-2-Raw"}, "minimax-h3": {"pretrained_model_name_or_path": "MiniMaxAI/MiniMax-H3"}, } @@ -807,6 +813,9 @@ def infer_diffusers_model_type(checkpoint): elif any(key in checkpoint for key in CHECKPOINT_KEY_NAMES["qwen-image-2.1"]): model_type = "qwen-image-2.1" + 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" @@ -4419,6 +4428,33 @@ def convert_qwen_image21_transformer_checkpoint_to_diffusers(checkpoint, **kwarg return converted_state_dict +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): prefix_rename_dict = { "first.": "img_in.", diff --git a/tests/single_file/test_model_ideogram4_transformer_single_file.py b/tests/single_file/test_model_ideogram4_transformer_single_file.py new file mode 100644 index 000000000000..e16c6ceeb521 --- /dev/null +++ b/tests/single_file/test_model_ideogram4_transformer_single_file.py @@ -0,0 +1,305 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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. + +""" +Tests for Ideogram4Transformer2DModel from_single_file() support. + +Two suites: + 1. TestIdeogram4ConversionFunction — offline unit tests for the key-mapping + function. No checkpoint, no GPU, no network. + 2. TestIdeogram4TransformerSingleFile — @nightly integration test that loads + an actual single-file checkpoint and compares against from_pretrained(). + Requires a public checkpoint; fill in ckpt_path / repo_id when available. +""" + +import unittest + +import torch + +from diffusers import Ideogram4Transformer2DModel +from diffusers.loaders.single_file_utils import ( + DIFFUSERS_DEFAULT_PIPELINE_PATHS, + convert_ideogram4_transformer_checkpoint_to_diffusers, + infer_diffusers_model_type, +) + +from ..testing_utils import enable_full_determinism + + +enable_full_determinism() + + +# --------------------------------------------------------------------------- +# Helper: build a synthetic "original" state dict for a small Ideogram4 model +# --------------------------------------------------------------------------- + + +def _make_ideogram4_original_state_dict( + num_layers=2, + hidden_size=32, + num_heads=4, + intermediate_size=32, + adaln_dim=16, + in_channels=16, + llm_features_dim=24, + prefix="diffusion_model.", +): + """ + Produces a synthetic state dict in the pre-diffusers Ideogram4 format: + - fused attention.qkv (3 * head_dim * num_heads, hidden_size) + - attention.o (hidden_size, hidden_size) + - feed_forward.w1/w2/w3 + - adaln_modulation + - per-layer norms + - top-level embeddings / final_layer + All tensors contain random data; shapes match what the model actually expects. + """ + head_dim = hidden_size // num_heads + sd = {} + + # Top-level modules + sd[f"{prefix}input_proj.weight"] = torch.randn(hidden_size, in_channels) + sd[f"{prefix}input_proj.bias"] = torch.randn(hidden_size) + sd[f"{prefix}llm_cond_norm.weight"] = torch.randn(llm_features_dim) + sd[f"{prefix}llm_cond_proj.weight"] = torch.randn(hidden_size, llm_features_dim) + sd[f"{prefix}llm_cond_proj.bias"] = torch.randn(hidden_size) + # t_embedding is Ideogram4EmbedScalar: two linear layers (mlp_in, mlp_out) both with bias + sd[f"{prefix}t_embedding.mlp_in.weight"] = torch.randn(hidden_size, hidden_size) + sd[f"{prefix}t_embedding.mlp_in.bias"] = torch.randn(hidden_size) + sd[f"{prefix}t_embedding.mlp_out.weight"] = torch.randn(hidden_size, hidden_size) + sd[f"{prefix}t_embedding.mlp_out.bias"] = torch.randn(hidden_size) + sd[f"{prefix}adaln_proj.weight"] = torch.randn(adaln_dim, hidden_size) + sd[f"{prefix}adaln_proj.bias"] = torch.randn(adaln_dim) + # embed_image_indicator is nn.Embedding (weight only, no bias) + sd[f"{prefix}embed_image_indicator.weight"] = torch.randn(2, hidden_size) + # rotary_emb uses register_buffer with persistent=False -> not in state_dict + + for i in range(num_layers): + p = f"{prefix}layers.{i}." + # Fused QKV (original format, no bias) + sd[f"{p}attention.qkv.weight"] = torch.randn(3 * hidden_size, hidden_size) + # Output projection (original name: attention.o, no bias) + sd[f"{p}attention.o.weight"] = torch.randn(hidden_size, hidden_size) + # q/k norms (RMSNorm with elementwise_affine=True, same names in both formats) + sd[f"{p}attention.norm_q.weight"] = torch.randn(head_dim) + sd[f"{p}attention.norm_k.weight"] = torch.randn(head_dim) + # Feed-forward (no bias) + sd[f"{p}feed_forward.w1.weight"] = torch.randn(intermediate_size, hidden_size) + sd[f"{p}feed_forward.w2.weight"] = torch.randn(hidden_size, intermediate_size) + sd[f"{p}feed_forward.w3.weight"] = torch.randn(intermediate_size, hidden_size) + # AdaLN modulation (has bias) + sd[f"{p}adaln_modulation.weight"] = torch.randn(4 * hidden_size, adaln_dim) + sd[f"{p}adaln_modulation.bias"] = torch.randn(4 * hidden_size) + # RMSNorms (same names in both formats) + for norm in ("attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2"): + sd[f"{p}{norm}.weight"] = torch.randn(hidden_size) + + # Final layer (same names in both formats) + # norm_final uses elementwise_affine=False -> no weight/bias in state_dict + sd[f"{prefix}final_layer.linear.weight"] = torch.randn(in_channels, hidden_size) + sd[f"{prefix}final_layer.linear.bias"] = torch.randn(in_channels) + # adaln_modulation: nn.Linear(adaln_dim, hidden_size, bias=True) + sd[f"{prefix}final_layer.adaln_modulation.weight"] = torch.randn(hidden_size, adaln_dim) + sd[f"{prefix}final_layer.adaln_modulation.bias"] = torch.randn(hidden_size) + + return sd + + +class TestIdeogram4ConversionFunction(unittest.TestCase): + """Offline unit tests for convert_ideogram4_transformer_checkpoint_to_diffusers.""" + + def _convert(self, prefix="diffusion_model.", num_layers=2): + sd = _make_ideogram4_original_state_dict(prefix=prefix, num_layers=num_layers) + return convert_ideogram4_transformer_checkpoint_to_diffusers(sd) + + # ------------------------------------------------------------------ + # Prefix stripping + # ------------------------------------------------------------------ + + def test_strips_diffusion_model_prefix(self): + result = self._convert(prefix="diffusion_model.") + self.assertFalse(any(k.startswith("diffusion_model.") for k in result)) + + def test_strips_conditional_transformer_prefix(self): + sd = _make_ideogram4_original_state_dict(prefix="conditional_transformer.") + result = convert_ideogram4_transformer_checkpoint_to_diffusers(sd) + self.assertFalse(any(k.startswith("conditional_transformer.") for k in result)) + + def test_no_prefix_passthrough(self): + # If there is no known prefix the keys should survive untouched. + sd = _make_ideogram4_original_state_dict(prefix="") + result = convert_ideogram4_transformer_checkpoint_to_diffusers(sd) + self.assertIn("layers.0.attention.to_q.weight", result) + + # ------------------------------------------------------------------ + # QKV split + # ------------------------------------------------------------------ + + def test_fused_qkv_is_split_into_three_projections(self): + result = self._convert() + for i in range(2): + self.assertIn(f"layers.{i}.attention.to_q.weight", result) + self.assertIn(f"layers.{i}.attention.to_k.weight", result) + self.assertIn(f"layers.{i}.attention.to_v.weight", result) + self.assertNotIn(f"layers.{i}.attention.qkv.weight", result) + + def test_split_shapes_are_equal_thirds(self): + result = self._convert() + # Each split should be hidden_size // 3 of the original fused dim (hidden_size=32 -> 32 each) + for proj in ("to_q", "to_k", "to_v"): + self.assertEqual(result[f"layers.0.attention.{proj}.weight"].shape, (32, 32)) + + def test_split_values_are_contiguous_chunks(self): + # Build a deterministic fused weight and verify the three slices. + fused = torch.arange(96, dtype=torch.float32).reshape(96, 1) + sd = {"diffusion_model.layers.0.attention.qkv.weight": fused} + result = convert_ideogram4_transformer_checkpoint_to_diffusers(sd) + expected_q = fused[:32] + expected_k = fused[32:64] + expected_v = fused[64:] + self.assertTrue(torch.equal(result["layers.0.attention.to_q.weight"], expected_q)) + self.assertTrue(torch.equal(result["layers.0.attention.to_k.weight"], expected_k)) + self.assertTrue(torch.equal(result["layers.0.attention.to_v.weight"], expected_v)) + + # ------------------------------------------------------------------ + # Output projection rename: attention.o -> attention.to_out.0 + # ------------------------------------------------------------------ + + def test_attention_o_renamed_to_to_out_0(self): + result = self._convert() + for i in range(2): + self.assertIn(f"layers.{i}.attention.to_out.0.weight", result) + self.assertNotIn(f"layers.{i}.attention.o.weight", result) + + # ------------------------------------------------------------------ + # Keys that should pass through unchanged + # ------------------------------------------------------------------ + + def test_feed_forward_keys_unchanged(self): + result = self._convert() + for i in range(2): + for w in ("w1", "w2", "w3"): + self.assertIn(f"layers.{i}.feed_forward.{w}.weight", result) + + def test_norm_keys_unchanged(self): + result = self._convert() + for i in range(2): + for norm in ("attention_norm1", "attention_norm2", "ffn_norm1", "ffn_norm2"): + self.assertIn(f"layers.{i}.{norm}.weight", result) + + def test_adaln_modulation_unchanged(self): + result = self._convert() + for i in range(2): + self.assertIn(f"layers.{i}.adaln_modulation.weight", result) + + def test_final_layer_keys_unchanged(self): + result = self._convert() + # norm_final uses elementwise_affine=False, so it has no parameters + self.assertNotIn("final_layer.norm_final.weight", result) + self.assertIn("final_layer.linear.weight", result) + self.assertIn("final_layer.linear.bias", result) + self.assertIn("final_layer.adaln_modulation.weight", result) + self.assertIn("final_layer.adaln_modulation.bias", result) + + def test_top_level_embedding_keys_unchanged(self): + result = self._convert() + for key in ( + "input_proj.weight", + "input_proj.bias", + "llm_cond_norm.weight", + "llm_cond_proj.weight", + "llm_cond_proj.bias", + # t_embedding is Ideogram4EmbedScalar: two linear layers + "t_embedding.mlp_in.weight", + "t_embedding.mlp_in.bias", + "t_embedding.mlp_out.weight", + "t_embedding.mlp_out.bias", + "adaln_proj.weight", + "adaln_proj.bias", + "embed_image_indicator.weight", + ): + self.assertIn(key, result) + + # ------------------------------------------------------------------ + # Round-trip key coverage: every diffusers state dict key is produced + # ------------------------------------------------------------------ + + def test_output_keys_match_diffusers_model_state_dict(self): + """ + Instantiate a tiny Ideogram4 model, read its diffusers-format state dict, + then verify that the conversion function produces exactly the same key set + when given a synthetic original-format checkpoint. + """ + init_cfg = { + "in_channels": 16, + "num_layers": 2, + "attention_head_dim": 8, + "num_attention_heads": 4, + "intermediate_size": 32, + "adaln_dim": 16, + "llm_features_dim": 24, + "rope_theta": 10_000, + "mrope_section": (2, 1, 1), + "norm_eps": 1e-5, + } + model = Ideogram4Transformer2DModel(**init_cfg) + diffusers_keys = set(model.state_dict().keys()) + + original_sd = _make_ideogram4_original_state_dict( + num_layers=2, + hidden_size=32, # num_heads * head_dim = 4 * 8 + num_heads=4, + intermediate_size=32, + adaln_dim=16, + in_channels=16, + llm_features_dim=24, + prefix="diffusion_model.", + ) + converted = convert_ideogram4_transformer_checkpoint_to_diffusers(original_sd) + converted_keys = set(converted.keys()) + + missing = diffusers_keys - converted_keys + extra = converted_keys - diffusers_keys + self.assertEqual(missing, set(), msg=f"Keys in diffusers model but missing from conversion: {missing}") + self.assertEqual(extra, set(), msg=f"Keys produced by conversion but not in diffusers model: {extra}") + + +class TestIdeogram4ModelTypeInference(unittest.TestCase): + """Offline unit tests for detecting Ideogram4 checkpoints in infer_diffusers_model_type.""" + + def test_infers_ideogram4_for_all_prefixes(self): + for prefix in ("diffusion_model.", "conditional_transformer.", ""): + with self.subTest(prefix=prefix): + sd = _make_ideogram4_original_state_dict(prefix=prefix) + self.assertEqual(infer_diffusers_model_type(sd), "ideogram4") + + def test_default_pipeline_path_registered(self): + self.assertIn("ideogram4", DIFFUSERS_DEFAULT_PIPELINE_PATHS) + + +# --------------------------------------------------------------------------- +# Integration test — requires actual single-file checkpoint (nightly / big GPU) +# --------------------------------------------------------------------------- + +# TODO: fill in once a public Ideogram4 single-file checkpoint is released. +# class TestIdeogram4TransformerSingleFile(SingleFileModelTesterMixin): +# model_class = Ideogram4Transformer2DModel +# ckpt_path = "https://huggingface.co///blob/main/ideogram4_transformer.safetensors" +# repo_id = "/" +# subfolder = "transformer" + + +if __name__ == "__main__": + unittest.main()