Skip to content

Z-Image-Turbo text-to-image support#24

Merged
TroyHernandez merged 10 commits into
mainfrom
feature/zimage
Jul 8, 2026
Merged

Z-Image-Turbo text-to-image support#24
TroyHernandez merged 10 commits into
mainfrom
feature/zimage

Conversation

@TroyHernandez

Copy link
Copy Markdown
Contributor

Adds Z-Image-Turbo (Tongyi-MAI, Apache-2.0) text-to-image: a 6B
guidance-distilled single-stream DiT that renders legible text in
images (English and Chinese), 8 steps, no CFG.

download_zimage_turbo()   # ungated; ~33 GB download, one-time fp8
                          # quantize to a 5.9 GB artifact
txt2img_zimage("A sign that reads \"DIFFUSER\" carved in wood",
               seed = 42)   # or txt2img(..., model_name = "zimage")

Architecture (ported from the diffusers reference, Apache-2.0)

  • Transformer: single-stream joint attention (no MMDiT split);
    30 trunk blocks + 2 modulated noise-refiner + 2 unmodulated
    context-refiner blocks; sandwich RMSNorms (learned norm before AND
    after attention/FF); scale/gate-only modulation with tanh-squashed
    gates (no shift); SwiGLU with separate gate weights; 3-axis RoPE
    theta 256; sub-sequences padded to multiples of 32 with learned pad
    tokens. flux_attention reused as-is (bias = FALSE).
  • Text encoder: the FLUX.2-klein Qwen3-4B port reused verbatim
    (tokenizer files are byte-identical between the checkpoints); only
    the chat template differs (enable_thinking = TRUE, no think block)
    and the pipeline takes the penultimate hidden state, mask-sliced to
    the caption length.
  • VAE: the FLUX.1 16-channel decoder unchanged (the checkpoint's
    VAE config is the flux-dev autoencoder).
  • Scheduler: FlowMatch static shift 3.0 on sigmas
    linspace(1, 1/N, N); the model consumes the reversed normalized
    timestep (1000 - t)/1000 and its output is negated before the Euler
    step.
  • Quantization: one-time local fp8 (or NF4) through the existing
    family-hook machinery; cast set is the 7 core linears x 34 blocks =
    238 weights, adaLN/embedders stay bf16. The fp32-shipped checkpoint's
    pass-through residents are stored bf16.

Allocator gc-gate fix (cross-model)

Generation initially ran at ~142 s regardless of resolution; profiling
showed 89% of wall time in R gc (~2,200 gcs x 62 ms). torch's
allocated/reserved ratio gate (0.95) is chronically exceeded under
backend:native (steady state 0.957), firing the R-gc callback on
nearly every CUDA allocation — and start_torch() reads the gate
options only once at init, so mid-session option changes are inert.
.flux_gc_gates() now pushes thresholds into the live allocator:
ratio gate off, absolute allocated_rate 0.65 as the
garbage-accumulation gate, footprint-based reserved_rate as the net.

Measured 1024x1024 on the RTX 5060 Ti 16 GB:

Model Before After
Z-Image-Turbo 143 s 24 s
FLUX.2 klein 48 s 13 s
FLUX.1-schnell 121 s 55 s

Verification

  • Per-module parity fixtures against the diffusers reference: position
    ids bit-exact, full tiny-model forward < 1e-4, strict bidirectional
    key censuses, chat-template token ids exact against the shipped
    tokenizer (CJK/emoji included), quantized round trips (NF4 + streamed
    and resident fp8). Full suite: 692/692.
  • AST-based structural audit of the translation (numeric literals +
    call names per paired scope): clean; only dead defaults and
    omni-only constants differ.
  • Full-size load census exact (238 cast weights); same-seed output
    reproduced identically before/after the gate fix.
  • R CMD check clean.

Sample outputs in ~/Sync (zimage_sign_1024.png shows the text
rendering).

…R gc

The torch allocated/reserved ratio gate (0.95) is chronically exceeded
under backend:native (measured 0.957 steady state), firing the R gc
callback on nearly every CUDA allocation: ~2,200 gcs x 62 ms per
Z-Image generation. start_torch() reads the gate options only once at
torch init, so the loaders now push thresholds into the live allocator
directly: ratio gate off, absolute allocated_rate 0.65 as the
garbage-accumulation gate, footprint-based reserved_rate as the net.

Measured 1024x1024 on the RTX 5060 Ti 16 GB:
  Z-Image-Turbo 143 -> 24 s, klein 48 -> 13 s, FLUX.1 121 -> 55 s
@TroyHernandez TroyHernandez merged commit 58433a0 into main Jul 8, 2026
0 of 2 checks passed
@TroyHernandez TroyHernandez deleted the feature/zimage branch July 8, 2026 20:30
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant