Skip to content

Conversation

@hlky
Copy link
Contributor

@hlky hlky commented Dec 25, 2025

What does this PR do?

  1. CHECKPOINT_KEY_NAMES changed from cap_embedder.0.weight to layers.0.adaLN_modulation.0.weight and model.diffusion_model.layers.0.adaLN_modulation.0.weight. This helps better distinguish vs lumina, and model.diffusion_model. variant is also required for some checkpoints.
  2. In convert_z_image_transformer_checkpoint_to_diffusers {"model.diffusion_model.": ""} is added to Z_IMAGE_KEYS_RENAME_DICT to remove the prefix
  3. Special case norm_final.weight key is removed, this appears to be an erroneous additional weight, we have norm_final as nn.LayerNorm in all_final_layer.2-1 (FinalLayer module) however elementwise_affine=False so weight is None, also if we check the weight itself in an example checkpoint (unstableRevolution_Bf16.safetensors) we see some nan tensor([1.8447e+19, 1.9886e+31, nan, ..., 7.3242e-02, 4.6587e-20, 1.3733e-02], dtype=torch.bfloat16) and other extreme essentially junk values.

Repro

import torch
from diffusers.models import ZImageTransformer2DModel
from huggingface_hub import hf_hub_download

model = ZImageTransformer2DModel.from_single_file(
    hf_hub_download(
        "hlky/unstable-revolution-zit", "unstableRevolution_Bf16.safetensors"
    ),
    torch_dtype=torch.bfloat16,
)

Fixes #12887

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

cc @vladmandic

@HuggingFaceDocBuilderDev

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.

@hlky
Copy link
Contributor Author

hlky commented Dec 29, 2025

Failing tests are unrelated, failure since Z-Image was introduced https://github.com/huggingface/diffusers/actions/runs/20073726260/job/57582516584 - however I can't reproduce locally, all fast tests succeed in my environment.

@DN6 DN6 merged commit 4737806 into huggingface:main Jan 2, 2026
10 of 11 checks passed
Aki-07 pushed a commit to bconstantine/diffusers that referenced this pull request Jan 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

z-image-turbo from_single_file errors

3 participants