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.

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

Labels

None yet

3 participants