main gained a config-driven Qwen3 text encoder (Qwen3ForCausalLM + Qwen3TextConfig)
via the Flux2 Klein port (#1349). Reuse it for Z-Image instead of the parallel
~540-line bespoke Qwen3Model this PR originally added: drop the duplicate encoder
+ config, and map Z-Image-Turbo's "Qwen3Model" architecture string to the shared
encoder in the registry. update_model_arch populates Z-Image's dims from config.json.
Validating the reuse on L40S surfaced a real, batch-only divergence in the shared
encoder: it built position_ids as [1, seq_len], but the rotary layer flattens
positions to num_tokens and reshapes q/k to (num_tokens, -1, head_dim). For
batch>1 that folded the batch dim into the head dim and misaligned RoPE (fp32
mean_diff 0.22 at batch=2; batch=1 was fine). Fix: expand position_ids to
[batch_size, seq_len]. batch=1 is byte-identical, so Flux2 Klein is unaffected;
this also fixes a latent batch bug in the shared encoder.
Encoder parity now PASSES on L40S (Z-Image-Turbo): fp32 bit-exact against the
shared encoder (both batch elements, max=0.0000); bf16 within the existing
thresholds (last_hidden mean ~0.016, pre-norm mean ~0.07-0.08).
Also in this PR (unchanged): the gated scheduler option
use_reference_discrete_timesteps (default False) for Z-Image timestep parity,
the Z-Image component parity tests (encoder/scheduler/tokenizer/VAE), and the
encoder parity test's strict-load allowlist + OOM-safe two-model handling
(free HF ref + gc.collect before loading FastVideo).
tests/local_tests/ holds local-only parity, smoke, and component tests
organized by model family (e.g. tests/local_tests/ltx2/,
tests/local_tests/stable_audio/). They typically require a checked-out
reference repo and/or model weights and are skipped in CI; see each family's
README.md for setup. The top-level
tests/local_tests/README.md is the family index.
The CI-backed test suite still lives in fastvideo/tests/.
Raghav
bf839dde18