[boshjerns]

Summary

ZImageCheckpointer.build_text_encoder reads rope_theta from the top level of the text encoder config dict. transformers 5.x moved it into rope_parameters, so with transformers 5.x loading Z-Image from the Diffusers checkpoint fails with KeyError: 'rope_theta'. This reads rope_parameters["rope_theta"] when it is there and falls back to the top-level key, the same way the FLUX.2-klein loader already does (generate_flux2klein.py).

Issue

pip install . takes its dependencies from dependencies/requirements/generated_requirements/requirements.txt, which allows transformers>=4.57.6, so a fresh install can resolve transformers 5.x (base_requirements says transformers<5.0.0). With transformers 5.20.1:

from transformers import AutoConfig
cfg = AutoConfig.from_pretrained("Tongyi-MAI/Z-Image-Turbo", subfolder="text_encoder").to_dict()
print(cfg.get("rope_theta"), cfg.get("rope_parameters"))
# None {'rope_theta': 1000000, 'rope_type': 'default'}

load_diffusers_pipeline passes this dict to build_text_encoder, which raises:

  File ".../checkpointing/z_image_checkpointer.py", line 241, in build_text_encoder
    rope_theta=text_encoder_config["rope_theta"],
KeyError: 'rope_theta'

The Orbax restore path calls build_text_encoder with the config saved in the checkpoint metadata, so restoring a cache whose metadata was written under transformers 5.x hits the same error (caches written under 4.x still restore).

Changes

• checkpointing/z_image_checkpointer.py: take rope_theta from rope_parameters if present, otherwise from the top level. If neither exists it still raises KeyError, as before.
• tests/z_image/z_image_checkpointer_test.py (new): builds the text encoder abstractly from a tiny Qwen3 config (no weights downloaded) and checks rope_theta for the 4.x layout, the 5.x layout, and AutoConfig.for_model("qwen3", ...).to_dict() from whatever transformers is installed, plus that a config with no rope_theta anywhere still raises KeyError.

Testing

New test, JAX CPU:

transformers	main	this PR
5.20.1	2 failed (KeyError: 'rope_theta'), 2 passed	4 passed
4.57.6	1 failed (5.x layout case), 3 passed	4 passed

Real weights with transformers 5.20.1: I built and loaded the Tongyi-MAI/Z-Image-Turbo text encoder through build_text_encoder and load_qwen3_weights in float32 and compared the layer encode_step returns (all_hidden_states[-2], index 35) with hidden_states[35] from the transformers PyTorch model in float32, on a 17-token prompt under jax.default_matmul_precision("highest"): per-token cosine 1.000000 (min and mean), relative L2 2.7e-6. On main the same script stops at the KeyError.

Checks: pyink --check --pyink-indentation=2 --line-length=125 on the touched files is clean, ruff check . from the repo root passes, pylint reports no new messages (the checkpointer file has the same 5 existing messages as on main, the test file rates 10.00).

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[boshjerns]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates the build_text_encoder method in z_image_checkpointer.py to support both transformers 4.x and 5.x configuration formats by checking for rope_theta within rope_parameters as well as at the top level of the configuration dictionary. Additionally, a new test suite z_image_checkpointer_test.py has been introduced to verify correct behavior across these different configuration structures. There are no review comments, and I have no further feedback to provide.