Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Add from_single_file support to ErnieImageTransformer2DModel#13727

Merged
DN6 merged 2 commits into
huggingface:mainfrom
akshan-main:ernie-image-from-single-file
Jun 8, 2026
Merged

Add from_single_file support to ErnieImageTransformer2DModel#13727
DN6 merged 2 commits into
huggingface:mainfrom
akshan-main:ernie-image-from-single-file

Conversation

@akshan-main

@akshan-main akshan-main commented May 12, 2026

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

@sayakpaul sayakpaul requested a review from DN6 May 12, 2026 07:38
wadeKeith
wadeKeith reviewed May 12, 2026
View reviewed changes

@wadeKeith wadeKeith left a comment •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Withdrawn; this review was too cursory.

@DN6

DN6 commented May 13, 2026

Copy link
Copy Markdown
Collaborator

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

@akshan-main

akshan-main commented May 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@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

@akshan-main

akshan-main commented May 22, 2026

Copy link
Copy Markdown
Contributor Author

Friendly ping @DN6

DN6
DN6 approved these changes May 27, 2026
View reviewed changes

@DN6 DN6 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Sorry for the delay. One minor change requested.

Comment thread src/diffusers/loaders/single_file_utils.py Outdated
Comment on lines +4175 to +4176
for k in list(checkpoint.keys()):
converted_state_dict[k] = checkpoint.pop(k)

@DN6 DN6 May 19, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

We don't really need this part?

@akshan-main akshan-main May 27, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Oh yeah! I followed the flux/chroma return shape (separate converted_state_dict) but ernie doesn't need any further remapping. Fixed!

@akshan-main

akshan-main commented May 27, 2026

Copy link
Copy Markdown
Contributor Author

@DN6 all good now!

@akshan-main akshan-main requested a review from DN6 May 27, 2026 13:40
@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Jun 4, 2026

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@akshan-main akshan-main requested a review from DN6 June 4, 2026 12:50
@DN6 DN6 merged commit 0a5c98e into huggingface:main Jun 8, 2026
15 checks passed
DN6 pushed a commit that referenced this pull request Jul 1, 2026
@akshan-main @DN6
Add from_single_file support to ErnieImageTransformer2DModel (#13727)
9dde5c3 
* Add from_single_file support to ErnieImageTransformer2DModel

* drop redundant copy loop
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@DN6 DN6 Awaiting requested review from DN6
1 more reviewer
@wadeKeith wadeKeith wadeKeith left review comments
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

fixes-issue models single-file size/S PR with diff < 50 LOC

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

ErnieImageTransformer2DModel does not implement from_single_file

4 participants

@akshan-main @DN6 @HuggingFaceDocBuilderDev @wadeKeith