Fix: resolve Qwen3 text encoder loading issues for FP8 and GGUF formats

[arifanchan]

Description

This PR resolves three sequential crashing issues that occur when loading and compiling the Qwen3 text encoder in different quantized formats (FP8 Safetensors and GGUF) alongside the Z-Image model.

Issue 1: Missing lm_head.weight in FP8 Safetensors

Error: Exception: Missing keys: ['lm_head.weight']

• Root Cause: FP8 checkpoints intentionally drop the lm_head.weight key to save disk space, causing mmgp/offload.py strict key checks to fail.
• Solution: Added _fix_qwen_fp8_sd preprocessor in z_image_main.py to dynamically alias the model.embed_tokens.weight memory reference into the lm_head.weight slot on the fly.

Issue 2: SDPA Dtype Mismatch in GGUF

Error: RuntimeError: Expected query, key, and value to have the same dtype

• Root Cause: llama.cpp integration often provides the Value tensor in bfloat16, while HF transformers computes RoPE for Query/Key in float32. PyTorch SDPA strictly requires uniformity.
• Solution: Added a .to(dtype) cast immediately after loading the text encoder in z_image_main.py to enforce strict dtype uniformity.

Issue 3: Dynamo RecursionError with Q8_0 GGUF

Error: torch._dynamo.exc.InternalTorchDynamoError: RecursionError: maximum recursion depth exceeded

• Root Cause: The Q8_0 GGUF format lacks an embedding fast path, causing it to fall back to a custom Python __torch_function__ subclass wrapper. When these embeddings are passed to the torch.compile wrapped Z-Image transformer, Dynamo tries to trace the custom wrapper and enters an infinite loop.
• Solution: Added an .as_subclass(torch.Tensor) cast to the embeddings in models/z_image/pipeline_z_image.py to strip the GGUF wrapper before they enter the compiled transformer. This has zero VRAM/performance penalty and acts as a safe bridge to eager mode.

[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!