Z-Image-Turbo from_single_file

[hlky]

What does this PR do?

importtorchfromdiffusersimportGGUFQuantizationConfigfromdiffusers.modelsimportZImageTransformer2DModelfromhuggingface_hubimporthf_hub_downloadmodel=ZImageTransformer2DModel.from_single_file(
hf_hub_download("jayn7/Z-Image-Turbo-GGUF", "z_image_turbo-Q3_K_S.gguf"),
quantization_config=GGUFQuantizationConfig(compute_dtype=torch.bfloat16),
torch_dtype=torch.bfloat16,
)
model=ZImageTransformer2DModel.from_single_file(
hf_hub_download(
"Comfy-Org/z_image_turbo",
"split_files/diffusion_models/z_image_turbo_bf16.safetensors",
)
)

See https://huggingface.co/Comfy-Org/z_image_turbo/blob/main/z_image_convert_original_to_comfy.py

Fixes#12748

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

[Vargol]

Sorry if this shows as a dupe I though commented this ages ago, but there's no sign of it
Hi, thanks for this, the model I tried loads but t doesn't work, if looks like there's some code in the transformer code that checks the dtype of the weights and gets the int8 storage dtype instead of the GGUF compute dtype and then calls (eventually) torch.nn.Linear with the wrong types.

The code is the TimestepEmbedder forward function

defforward(self, t):
t_freq=self.timestep_embedding(t, self.frequency_embedding_size)
weight_dtype=self.mlp[0].weight.dtypeifweight_dtype.is_floating_point:
t_freq=t_freq.to(weight_dtype)
t_emb=self.mlp(t_freq)
returnt_emb

self.mlp[0].weight.dtype returns int8 for a GGUF format model

This leads to the forward code

t_emb=self.mlp(t_freq)

eventually calling output = torch.nn.functional.linear(inputs, weight, bias) with torch.Float , torch.BFloat16, torch.BFloat16 arguments.
on MPS this fails with Destination NDArray and Accumulator NDArray cannot have different datatype in MPSNDArrayMatrixMultiplication , I don't have any CUDA or other devices to check if this is a generic
issue.

If I hardcode the right type, I can generate an image without issue

t_freq=t_freq.to(self.mlp[0].compute_dtype)

Presumably It will need to incorporated properly with a attribute check for compute_dtype as part of the dtype setting code rather than my brute force method

[hlky]

@Vargol Apologies, I only tested loading. Something like c84e6d7 should work

[Vargol]

That looks like it'll work, I'll give a quick test.

[Vargol]

Yep - that's worked, no errors only images :-)

[DN6]

Just a question. Why the device cast here? Is it to fix something else?

[hlky]

Oh, I meant to remove that, for context this patch was shared in the community to fix layer offloading in one of the training UIs, I was just curious what changes they made and forgot to revert before I started this branch, not sure if it's related to Diffusers offloading or specific to the third party repo. Removed in da06a2c

[DN6]

Just a question. Why the device cast here? Is it to fix something else?

[DN6]

Thanks @hlky 👍🏽

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.