Skip to content

z-image: ConvRot safetensors loader - #2893

Merged
bghira merged 1 commit into
mainfrom
feature/zimageturbo-convrot
Jul 25, 2026
Merged

z-image: ConvRot safetensors loader#2893
bghira merged 1 commit into
mainfrom
feature/zimageturbo-convrot

Conversation

@bghira

@bghira bghira commented Jul 25, 2026

Copy link
Copy Markdown
Owner

This pull request introduces support for loading quantized Z-Image ConvRot checkpoints using the SDNQ format and adds a new single-file loader interface to ZImageTransformer2DModel. It also includes comprehensive utilities for mapping, validating, and loading quantized weights, as well as new tests to verify the mapping logic and loader interface.

Quantized checkpoint loading and mapping:

  • Added quantized_loading.py, which implements the load_zimage_comfy_convrot_checkpoint function and associated helpers to load Z-Image ConvRot checkpoints in the SDNQ format, including key mapping, quantization metadata decoding, and module/buffer assignment. This enables efficient loading and usage of quantized models.
  • Introduced mapping logic to convert ComfyUI-style checkpoint keys to the expected Diffusers model keys, and validation for quantization metadata and tensor shapes.

Model API enhancements:

  • Added a from_single_file classmethod to ZImageTransformer2DModel for convenient instantiation from a single-file checkpoint, delegating to the new quantized loader.
  • Updated ZImage model logic to defer loading assistant LoRA weights until after quantization if the model is in SDNQ mode, ensuring correct initialization order.

Testing improvements:

  • Added unittests for the new key mapping, quantization metadata decoding, and the from_single_file loader interface, ensuring correct behavior and API coverage.

@bghira
bghira merged commit ee0beb5 into main Jul 25, 2026
2 checks passed
@bghira
bghira deleted the feature/zimageturbo-convrot branch July 25, 2026 14:53
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.

1 participant