Skip to content

Add from_single_file support to ErnieImageTransformer2DModel#13727

Open
akshan-main wants to merge 1 commit into
huggingface:mainfrom
akshan-main:ernie-image-from-single-file
Open

Add from_single_file support to ErnieImageTransformer2DModel#13727
akshan-main wants to merge 1 commit into
huggingface:mainfrom
akshan-main:ernie-image-from-single-file

Conversation

@akshan-main
Copy link
Copy Markdown
Contributor

What does this PR do?

Fixes #13722. ErnieImageTransformer2DModel does not implement from_single_file, unlike other models (Flux, Chroma, Z-Image, etc.), so loading a single-file checkpoint (e.g. a OneTrainer finetune) errors with AttributeError.

Adding standard FromOriginalModelMixin to the class and register a minimal checkpoint converter that strips the common model.diffusion_model. prefix(same shape as the Flux/Chroma converters)
Fixes #13722

Before submitting

Who can review?

@DN6 @sayakpaul

Copy link
Copy Markdown

@wadeKeith wadeKeith left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clean addition - from_single_file support for ErnieImageTransformer2DModel follows existing patterns. LGTM! Reviewed by Hermes Agent.

@DN6
Copy link
Copy Markdown
Collaborator

DN6 commented May 13, 2026

@akshan-main do you have an example snippet to verify that we can load the model?

@akshan-main
Copy link
Copy Markdown
Contributor Author

akshan-main commented May 13, 2026

@DN6 there's no public single-file ernie checkpoint yet, so I verified by saving the sharded baidu/ERNIE-Image transformer to a single safetensors file and then loading it back via from_single_file. Colab output with the PR's patches applied on top of main: #13722 (comment)

Output:

loading sharded checkpoint...
loaded, params: 8033490048
saved: /content/ernie_transformer_single.safetensors
loading via from_single_file...
loaded, params: 8033490048

key match: True
sample tensor 'x_embedder.proj.weight' equal: True

ROUND-TRIP OK

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ErnieImageTransformer2DModel does not implement from_single_file

3 participants