ErnieImageTransformer2DModel does not implement from_single_file

[dxqb]

Describe the bug

ErnieImageTransformer2DModel does not implement from_single_file, unlike other transformer models such as FluxTransformer2DModel, ChromaTransformer2DModel, and ZImageTransformer2DModel.

Reproduction

from diffusers import ErnieImageTransformer2DModel

transformer = ErnieImageTransformer2DModel.from_single_file(
"path/to/transformer.safetensors",
config="path/to/base_model",
subfolder="transformer",
)

Logs

Error:
  AttributeError: type object 'ErnieImageTransformer2DModel' has no attribute 'from_single_file'

System Info

tested with 0f1abc4 but it appears it was not added after this commit

Who can help?

Questions on single file checkpoints: @DN6

[akshan-main]

The class still inherits (ModelMixin, ConfigMixin, PeftAdapterMixin) without FromOriginalModelMixin, matching Flux/Chroma's pattern will fix this. Can I take it? @sayakpaul @DN6

[akshan-main]

tested on colab using

!rm -rf /content/diffusers
!git clone --depth=1 https://github.com/huggingface/diffusers.git /content/diffusers
!pip uninstall -y diffusers -q
!pip install -e /content/diffusers -q
!pip install -q transformers accelerate safetensors

# apply two-file patch
import pathlib

# Patch 1: transformer class — add FromOriginalModelMixin
tp = pathlib.Path("/content/diffusers/src/diffusers/models/transformers/transformer_ernie_image.py")
src = tp.read_text()
OLD_IMPORT = "from ..modeling_utils import ModelMixin"
NEW_IMPORT = "from ...loaders import FromOriginalModelMixin\nfrom ..modeling_utils import ModelMixin"
OLD_CLASS = "class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin):"
NEW_CLASS = "class ErnieImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin):"
assert OLD_IMPORT in src and OLD_CLASS in src
tp.write_text(src.replace(OLD_IMPORT, NEW_IMPORT).replace(OLD_CLASS, NEW_CLASS))

# Patch 2: add a minimal converter in single_file_utils.py
sfu = pathlib.Path("/content/diffusers/src/diffusers/loaders/single_file_utils.py")
src = sfu.read_text()
CONVERTER = '''

def convert_ernie_image_transformer_checkpoint_to_diffusers(checkpoint, **kwargs):
    converted = {}
    for k in list(checkpoint.keys()):
        new_k = k.replace("model.diffusion_model.", "")
        converted[new_k] = checkpoint.pop(k)
    return converted
'''
assert "convert_ernie_image_transformer_checkpoint_to_diffusers" not in src
sfu.write_text(src + CONVERTER)

# Patch 3: register in single_file_model.py
sf = pathlib.Path("/content/diffusers/src/diffusers/loaders/single_file_model.py")
src = sf.read_text()
OLD_IMPORTS = "from .single_file_utils import ("
NEW_IMPORTS = OLD_IMPORTS + "\n    convert_ernie_image_transformer_checkpoint_to_diffusers,"
assert OLD_IMPORTS in src and "convert_ernie_image_transformer_checkpoint_to_diffusers" not in src
src = src.replace(OLD_IMPORTS, NEW_IMPORTS)
OLD_REG = '''    "ChromaTransformer2DModel": {
        "checkpoint_mapping_fn": convert_chroma_transformer_checkpoint_to_diffusers,
        "default_subfolder": "transformer",
    },'''
NEW_REG = OLD_REG + '''
    "ErnieImageTransformer2DModel": {
        "checkpoint_mapping_fn": convert_ernie_image_transformer_checkpoint_to_diffusers,
        "default_subfolder": "transformer",
    },'''
assert OLD_REG in src
src = src.replace(OLD_REG, NEW_REG)
sf.write_text(src)
print("patched (3 files)")


# force-load patched diffusers, run round-trip test
import sys, torch
sys.path.insert(0, "/content/diffusers/src")
for m in list(sys.modules):
    if m == "diffusers" or m.startswith("diffusers."):
        del sys.modules[m]

from diffusers import ErnieImageTransformer2DModel
from safetensors.torch import save_file

print("loading sharded checkpoint...")
ref = ErnieImageTransformer2DModel.from_pretrained(
    "baidu/ERNIE-Image", subfolder="transformer", torch_dtype=torch.bfloat16
)
print("loaded, params:", sum(p.numel() for p in ref.parameters()))

single_path = "/content/ernie_transformer_single.safetensors"
save_file(ref.state_dict(), single_path)
print(f"saved: {single_path}")

print("loading via from_single_file...")
loaded = ErnieImageTransformer2DModel.from_single_file(
    single_path,
    config="baidu/ERNIE-Image",
    subfolder="transformer",
    torch_dtype=torch.bfloat16,
)
print("loaded, params:", sum(p.numel() for p in loaded.parameters()))

ref_sd, loaded_sd = ref.state_dict(), loaded.state_dict()
keys_match = set(ref_sd.keys()) == set(loaded_sd.keys())
sample_key = next(iter(ref_sd))
tensors_equal = torch.equal(ref_sd[sample_key], loaded_sd[sample_key])
print(f"\nkey match: {keys_match}")
print(f"sample tensor '{sample_key}' equal: {tensors_equal}")
print("\nROUND-TRIP OK" if (keys_match and tensors_equal) else "MISMATCH - converter needed")

output: ROUND-TRIP OK

can apply the same fix

[sayakpaul]

Feel free to take it! Thank you.