Fix: Resolve missing 'lm_head.weight' during Qwen3 FP8 text encoder loading

[arifanchan]

This PR fixes a crash that occurs in the queue worker when attempting to load Qwen3 FP8 text encoder safetensor checkpoints alongside the Z-Image Turbo 6B model.

Root Cause
Many quantized FP8 checkpoints drop the lm_head.weight key to save disk space because it is tied to and mathematically identical to model.embed_tokens.weight. When z_image_main.py calls offload.fast_load_transformers_model, the underlying memory manager (mmgp/offload.py) performs a strict key check and raises an exception: Exception: Missing keys: ['lm_head.weight'].

Solution
Rather than modifying the core mmgp/offload.py library to ignore missing keys (which could hide legitimate loading errors), this PR resolves the issue locally inside z_image_main.py:

• Implemented a fix_qwen_fp8_sd helper function.
• Passed this function to the preprocess_sd callback in fast_load_transformers_model.
• The callback intercepts the state dictionary as it loads from disk and dynamically duplicates model.embed_tokens.weight into the lm_head.weight slot before the strict key check occurs.

Changes Made
Updated models/z_image/z_image_main.py to inject the missing tied embedding weights on the fly during text encoder initialization.

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Adds a state-dict preprocessing hook when loading the Qwen3 text encoder to handle FP8 checkpoints that omit lm_head.weight.

Changes:

• Replaces direct fast_load_transformers_model call with a wrapper that injects missing lm_head.weight.
• Introduces fix_qwen_fp8_sd to patch the loaded state dict before model instantiation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[arifanchan]

Good catch on the semantics. We specifically want the tensor alias here rather than a .clone() to preserve the tied weights and avoid re-allocating the vocab matrix in VRAM. Updated the comment to clarify.

Good call on keeping the scope clean. Even though the loading script only runs once per model initialization, moving it to the module level as a private helper _fix_qwen_fp8_sd is tidier. Updated!