Error loading ZImageTurbo checkpoint

[jkoelker]

When loading the model from https://civitai.com/models/153568 we're getting:

❌ Failed to load model: Failed to load Gemma2Model. Weights for this component appear to be missing in the checkpoint.
Please load the component before passing it in as an argument to from_single_file.

text_encoder = Gemma2Model.from_pretrained('...')
pipe = StableDiffusionXLPipeline.from_single_file(<checkpoint path>, text_encoder=text_encoder)

[jkoelker]

@dobbyphus please research and analyze this issue

[dobbyphus]

Hey @jkoelker! I'm on it - researching this ZImageTurbo checkpoint loading issue.

I'll investigate:

1. The civitai model at https://civitai.com/models/153568
2. How our civitai_checkpoint.py pipeline handles loading
3. The Gemma2Model text encoder requirement
4. Potential solutions from diffusers library patterns

[dobbyphus]

Analysis Complete: ZImageTurbo Checkpoint Loading Issue

Summary

This is a known issue in the diffusers library, not a problem with Oneiro's code. The error occurs when loading certain ZImageTurbo checkpoints from CivitAI using from_single_file().

Root Cause

Two separate issues are at play:

Issue 1: Misidentified Error Message

The error message mentions Gemma2Model text encoder, but this is misleading. ZImageTurbo uses a Gemma2-based text encoder (from the Tongyi-MAI architecture), but the actual problem isn't related to loading the text encoder - it's a transformer model loading issue.

Issue 2: from_single_file Bug for ZImage Finetunes

This is tracked in huggingface/diffusers#12887. The bug affects most CivitAI finetunes of ZImageTurbo (not the base model).

Technical details from the diffusers issue:

• Many ZImage finetunes contain empty blocks that cause "meta tensor" errors
• The CHECKPOINT_KEY_NAMES detection was failing for some checkpoints
• Some checkpoints have a model.diffusion_model. prefix that wasn't being handled
• A spurious norm_final.weight key with junk/NaN values was causing issues

Fix Status

PR huggingface/diffusers#12888 has been approved and fixes these issues:

1. Improved checkpoint key detection
2. Handles model.diffusion_model. prefix in checkpoint keys
3. Removes the problematic norm_final.weight key

Recommendations for Oneiro

Short-term Workaround

The issue reports that using quantization_config works around the problem, as the quantization engine handles empty blocks properly.

Medium-term Solution

Once PR #12888 is merged, update to the latest diffusers version.

Long-term Consideration

The CivitaiCheckpointPipeline in Oneiro uses from_single_file() at line 630 in civitai_checkpoint.py:

self.pipe = pipeline_class.from_single_file(
    str(checkpoint_path),
    torch_dtype=self.policy.dtype,
)

No changes needed in Oneiro's code - this is purely a diffusers library bug.

Model-Specific Context

The model at civitai.com/models/153568 is a ZImageTurbo finetune. ZImageTurbo is:

• A 6B parameter image generation model from Tongyi-MAI (Alibaba)
• Uses a single-stream DiT architecture (S3-DiT)
• Employs Gemma2 as the text encoder
• Designed for 8-step inference

CivitAI labels this as compatible with "ZImagePipeline" in diffusers.

Action Items

1. Wait for diffusers fix: Track PR #12888 - it's approved and should be merged soon
2. Update diffusers: Once merged, update the diffusers dependency
3. No Oneiro code changes required: The fix is entirely in the upstream library

Let me know if you'd like me to:

• Monitor the PR and notify when it's merged
• Implement a temporary workaround using quantization config
• Create an issue to track the diffusers version update