[BUG?] Z Image Base safetensors model load errors - tensors wrong shape

[alex-ie]

Describe the Issue
I am trying to investigate leejet/stable-diffusion.cpp#1488 ([Bug] Z-Image Base GGUF checkpoints produce blank white images in stable-diffusion.cpp CUDA builds #1488) which I have experienced using kcpp 1.120.

I have downloaded original weights to test with them : https://huggingface.co/Tongyi-MAI/Z-Image/resolve/main/transformer/diffusion_pytorch_model-00001-of-00002.safetensors.

Now at model load time I get errors,r e.g. context_refiner.0.attention.qkv.weight ...got [3840,3840,1,1] expected [3840,4096,1,1]

This is my first time trying to load safetensors as main SD model, I assumed if vae / clip can be safetensors for kcpp, main model can be too and launcher allow to select it.

Back to ZiB: leejet quant has (per --analyze) context_refiner.0.attention.qkv.weight [3840,11520,1,1] and tensors for original safetensors file seems not to have context_refiner.0.attention.qkv at all, naming for some tensors is different. I do not know what to check further.

What could be the issue here? Do you need additional info? TIA

Additional Information:
1.120 Linux

[wbruna]

Split diffusers tensors (note the numbering 00001 of 00002) need a separate mapping (in this case https://huggingface.co/Tongyi-MAI/Z-Image/blob/main/transformer/diffusion_pytorch_model.safetensors.index.json) to determine which tensors are on which files. You typically need the whole diffusers model directory for loading them.

If you just want to compare with an unquantized version, my suggestion would be getting the ComfyUI .safetensors single file (https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/z_image.md has a link), which would be easier to work with. (note it's 'single file' in the sense that it contains the whole image model: you'll still need the llm and VAE).

[LostRuins]

Yeah please use single files for image gen models, split file loading probably doesnt work well

[alex-ie]

@wbruna, do you know if upstream https://github.com/leejet/stable-diffusion.cpp also cannot load split safetensors?

Thanks for (https://github.com/leejet/stable-diffusion.cpp/blob/master/docs/z_image.md) link, I have found the model file. There is also other file which looks like safetensors quant, maybe it is worth to test that too?

[wbruna]

I know upstream can load older/smaller models in diffusers format; it probably supports Z-Image too, but I've never tested it.

I've never used nvfp4, but it may be worth testing. And beware that int8_convrot is technically supported by both Koboldcpp and upstream, but the implementations are currently totally distinct.

[henk717]

I'd assume the model loading support is the same for them.
In general for the best results you need to load with gguf otherwise to my knowledge it becomes fp16 in memory unless its quantized.

[alex-ie]

Edit: I beg your pardon for being too rush as the model loaded using CPU backend, I guess due to BF16 larger, my 4GB VRAM is too small even with restricting VRAM usage to 1000 setting in launcher. Any way to split the model even more to use GPU? Or the cause of difference of load success might be elsewhere? TIA

Original: Both models from the link I found using link suggested by wbruna (https://huggingface.co/Comfy-Org/z_image/tree/main/split_files/diffusion_models) failed to load in kcpp right away with "Load Image Model OK:False", whereas original split model at least started to load (and have failed on wrong tensors shapes).