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

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

1 participant