Z-Image loader (LoadZImageTransformerModel) can't load ComfyUI-native checkpoints with a `model.diffusion_model.` key prefix

Summary

LoadZImageTransformerModel.loadmodel in comfyui/z_image/nodes.py already contains a convert_state_dict that maps ComfyUI-format Z-Image weights to the internal ZImageTransformer2DModel layout. But both the conversion gate and the converter assume bare keys. Many ComfyUI-native diffusion_models/*.safetensors checkpoints — e.g. CivitAI Z-Image fine-tunes like cyberrealisticZImage_v70_bf16.safetensors — instead store every tensor under a model.diffusion_model. prefix. For those, the gate is never triggered, the raw prefixed dict hits the strict load_state_dict, and loading fails:

RuntimeError: Error(s) in loading state_dict for ZImageTransformer2DModel:
  Missing key(s): x_pad_token, noise_refiner.0.attention.to_q.weight, ...
  Unexpected key(s): model.diffusion_model.cap_pad_token,
                     model.diffusion_model.context_refiner.0.attention.qkv.weight, ...

Root cause

In loadmodel:

transformer_state_dict = load_torch_file(model_path, safe_load=True)
def convert_state_dict(old_state_dict):
    ...
    if key.startswith('x_embedder.'): ...        # assumes bare keys
    elif '.attention.' in key: ...               # q_norm->norm_q, out->to_out.0, qkv-split
    ...
if "x_embedder.weight" in transformer_state_dict.keys():   # gate assumes bare key
    transformer_state_dict = convert_state_dict(transformer_state_dict)

When the checkpoint keys are model.diffusion_model.x_embedder.weight (etc.), "x_embedder.weight" in keys() is False, so convert_state_dict never runs, and the prefixed/fused-qkv dict is passed straight to transformer.load_state_dict(...) (strict) → failure.

Reproduction

1. Load any ComfyUI-native Z-Image checkpoint whose keys are prefixed with model.diffusion_model. (e.g. the CivitAI cyberrealisticZImage_v70_bf16.safetensors) via the VideoX-Fun LoadZImageTransformerModel node.
2. Observe the state_dict mismatch above.

The same weights, when saved without the wrapper prefix, load fine — confirming the prefix is the trigger, not an architecture difference.

Proposed fix (a few lines)

Strip the model.diffusion_model. prefix immediately after load, before the existing gate/converter — so the built-in conversion path handles the rest unchanged:

transformer_state_dict = load_torch_file(model_path, safe_load=True)

# Strip ComfyUI 'model.diffusion_model.' wrapper if present (e.g. CivitAI
# Z-Image fine-tunes). The converter + gate below assume bare keys.
_pref = "model.diffusion_model."
if any(k.startswith(_pref) for k in transformer_state_dict):
    transformer_state_dict = {
        (k[len(_pref):] if k.startswith(_pref) else k): v
        for k, v in transformer_state_dict.items()
    }

Verification

For cyberrealisticZImage_v70_bf16.safetensors (453 tensors, all model.diffusion_model.-prefixed):

• Before: strict load_state_dict fails (all expected keys missing, all checkpoint keys unexpected).
• After (strip + existing converter): exact match — 521 model keys == 521 produced, 0 missing, 0 extra — and strict load_state_dict succeeds with the real weights. (521 vs 453 because each fused qkv splits into three.)

The prefix strip is a no-op for reference-packaged / already-bare checkpoints, so existing models continue to load unchanged.

Environment

• ComfyUI + VideoX-Fun comfyui/ integration, LoadZImageTransformerModel node
• Example checkpoint: CivitAI cyberrealisticZImage_v70_bf16.safetensors (bf16, model.diffusion_model.-prefixed)
• Platform/torch-version independent (pure key-naming issue)

Companion report: this is one of two independent ComfyUI-integration bugs I hit; the other is filed separately.

Start in comfyui/z_image/nodes.py at LoadZImageTransformerModel.loadmodel and read the existing convert_state_dict gate. Reproduce with a ComfyUI-native checkpoint such as cyberrealisticZImage_v70_bf16.safetensors, then verify that prefixed weights load strictly with 521 model keys and that already-bare checkpoints remain unchanged.