Z-Image-Turbo text-to-image support

CLAUDE.md

@@ -318,6 +318,7 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt
 ### Model Support
 - [x] Add FLUX model support (FLUX.1-schnell, see below)
 - [x] Add FLUX.2 support (klein-4B, see below)
+- [x] Add Z-Image-Turbo support (text rendering, see below)
 - [ ] Add SD3 model support
 - [ ] ControlNet integration
 
@@ -336,8 +337,9 @@ txt2img_flux("An astronaut riding a horse on Mars, photorealistic",
 ```
 
 Measured on the RTX 5060 Ti 16 GB (NF4 resident, T5 float32 on CPU):
-1024x1024 in ~2 min wall (peak 8.7 GB alloc / 9.0 GB reserved),
-512x512 in ~1.5 min (peak 8.0 GB). CPU-only works too (~9 min at 256px).
+1024x1024 in ~55 s wall (peak 9.6 GB alloc / 9.8 GB reserved) after the
+allocator gc-gate fix (was ~2 min). CPU-only works too (~9 min at
+256px).
 
 Key components: `flux_transformer` (19 double + 38 single blocks),
 `t5_encoder` + `unigram_tokenizer` (pure R SentencePiece Viterbi),
@@ -363,8 +365,9 @@ txt2img_flux2("An astronaut riding a horse on Mars, photorealistic",
 ```
 
 Measured on the RTX 5060 Ti 16 GB (fp8 GPU-resident, Qwen3 bf16
-phase-onloaded): 1024x1024 in ~48 s (peak 8.2 GB), 512x512 in ~40 s;
-pipeline load 31 s. Cast census is exactly 104 weights.
+phase-onloaded): 1024x1024 in ~13 s (peak 12.5 GB alloc) after the
+allocator gc-gate fix (was ~48 s); pipeline load 31 s. Cast census is
+exactly 104 weights.
 
 Perf lesson that cost an afternoon: generation was 93.8% R garbage
 collection until the tokenizer stopped holding 151k-binding
@@ -373,6 +376,43 @@ and torch's allocator callbacks run gc hundreds of times per
 generation. Keep big lookup tables as atomic vectors (integer-id BPE
 via findInterval), never as environments.
 
+### Z-Image-Turbo (Complete)
+
+Guidance-distilled text-to-image (8 steps, no CFG), strong at legible
+text rendering (EN + CN). 6B single-stream DiT (sandwich RMSNorms,
+scale/gate-only tanh modulation, noise/context refiner stacks, 3-axis
+RoPE theta 256) + Qwen3-4B (thinking-enabled chat template, penultimate
+hidden state, mask-sliced captions) + the FLUX.1 16-channel VAE
+verbatim. FlowMatch with static shift 3.0; the model consumes the
+REVERSED timestep (1000 - t)/1000 and its output is negated.
+
+```r
+download_zimage_turbo()  # ungated, Apache-2.0; ~33 GB download (fp32
+                         # shards), one-time fp8 quantize to 5.9 GB
+txt2img_zimage("A sign that reads \"DIFFUSER\" carved in wood",
+               seed = 42)  # or txt2img("...", model_name = "zimage")
+```
+
+Measured on the RTX 5060 Ti 16 GB (fp8 GPU-resident, Qwen3 bf16
+phase-onloaded): 1024x1024 in ~24 s (peak 13.1 GB alloc / 13.9 GB
+reserved), 512x512 in ~12 s; pipeline load 41 s. Cast census is exactly
+238 weights (7 core linears x 34 blocks; adaLN + embedders stay bf16).
+
+Perf lesson (round two of the GC storm): torch's allocated/reserved
+ratio gate (0.95) is chronically exceeded under backend:native
+(steady-state ratio 0.957), so the R-gc callback fired on nearly every
+CUDA allocation — 89% of wall time (~2,200 gcs x 62 ms). And
+`start_torch()` reads the gate options only ONCE at torch init, so
+setting them mid-session is inert: push via
+`cpp_set_cuda_allocator_allocator_thresholds`. See `.flux_gc_gates()`.
+Fixed 1024x1024 across models: Z-Image 143->24 s, klein 48->13 s,
+FLUX.1 121->55 s.
+
+Gotchas that bit once: the caption's RoPE ramp is built over the PADDED
+length (pads continue the ramp; the (0,0,0) pad ids in the reference
+are dead code), and the checkpoint scheduler uses static shift 3.0 —
+the pipeline's calculate_shift/mu path is dead for this model.
+
 ### LTX-2.3 Video Generation (clean-room rewrite in progress)
 
 The original LTX-2.0 port was removed and is being replaced by a ground-up

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0.2
+Version: 0.1.0.3
 Authors@R: c(
     person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
            comment = c(ORCID = "0009-0005-4248-604X")),

NAMESPACE

@@ -13,6 +13,7 @@ export(download_flux1)
 export(download_flux2_klein)
 export(download_ltx2)
 export(download_model)
+export(download_zimage_turbo)
 export(encode_bpe)
 export(encode_qwen)
 export(encode_unigram)
@@ -183,6 +184,7 @@ export(txt2img_flux)
 export(txt2img_flux2)
 export(txt2img_sd21)
 export(txt2img_sdxl)
+export(txt2img_zimage)
 export(txt2vid_ltx2)
 export(unet_native)
 export(unet_native_from_torchscript)
@@ -193,6 +195,18 @@ export(vae_decoder_native)
 export(vocab_size)
 export(vram_report)
 export(write_wav)
+export(zimage_block)
+export(zimage_cap_pos_ids)
+export(zimage_feed_forward)
+export(zimage_final_layer)
+export(zimage_img_pos_ids)
+export(zimage_is_quant_key)
+export(zimage_load_pipeline)
+export(zimage_patchify)
+export(zimage_pos_embed)
+export(zimage_t_embedder)
+export(zimage_transformer)
+export(zimage_unpatchify)
 
 S3method(print,bpe_tokenizer)
 S3method(print,ltx23_checkpoint)

R/checkpoint_flux.R

@@ -61,6 +61,29 @@ flux_is_quant_key <- function(key) {
                                     ")\\.weight$"
 )
 
+# Z-Image cast set: the seven core linears (attention q/k/v/out and the
+# three SwiGLU weights) in all three block stacks. The per-block adaLN
+# modulation linears, embedders, pad tokens and norms stay in the
+# resident dtype. Full Turbo census: (30 layers + 2 noise_refiner + 2
+# context_refiner) x 7 = 238 cast weights, ~6.0B of the 6B parameters.
+.zimage_quant_cast_pattern <- paste0(
+                                     "^(noise_refiner|context_refiner|layers)\\.[0-9]+\\.(",
+                                     "attention\\.(to_q|to_k|to_v|to_out\\.0)",
+                                     "|feed_forward\\.(w1|w2|w3)",
+                                     ")\\.weight$"
+)
+
+#' Test whether a Z-Image key is in the quantization cast set
+#'
+#' @param key Character vector of parameter names (diffusers-style).
+#'
+#' @return Logical vector.
+#'
+#' @export
+zimage_is_quant_key <- function(key) {
+    grepl(.zimage_quant_cast_pattern, key)
+}
+
 #' Test whether a FLUX.2 key is in the quantization cast set
 #'
 #' @param key Character vector of parameter names (diffusers-style).
@@ -77,6 +100,8 @@ flux2_is_quant_key <- function(key) {
     cls <- config$`_class_name` %||% "FluxTransformer2DModel"
     if (identical(cls, "Flux2Transformer2DModel")) {
         "flux2"
+    } else if (identical(cls, "ZImageTransformer2DModel")) {
+        "zimage"
     } else {
         "flux1"
     }

R/dit_zimage.R

@@ -0,0 +1,149 @@
+#' Z-Image Transformer
+#'
+#' Fresh R port of ZImageTransformer2DModel from the diffusers reference
+#' (Apache-2.0, src/diffusers/models/transformers/transformer_z_image.py).
+#' Single-stream DiT: image tokens pass through a modulated noise
+#' refiner, caption tokens through an unmodulated context refiner, then
+#' both are concatenated (image first) and run through the main trunk.
+#' The module tree mirrors the reference state-dict keys 1:1
+#' (all_x_embedder.2-1, noise_refiner.N, context_refiner.N, layers.N,
+#' all_final_layer.2-1, t_embedder, cap_embedder, x_pad_token,
+#' cap_pad_token).
+#'
+#' This port is batch-of-1: \code{x} is a single latent [C, F, H, W] and
+#' \code{cap_feats} a single caption [L, cap_feat_dim], so sub-sequences
+#' are uniform and no attention mask is needed. Padding to a multiple of
+#' 32 tokens uses the learned pad parameters, appended after embedding
+#' (the reference pads raw features with repeats, embeds pointwise, then
+#' overwrites the pad rows with the same learned tokens).
+#'
+#' @param in_channels Integer. Latent channels. Default 16.
+#' @param dim Integer. Model width. Default 3840.
+#' @param n_layers Integer. Main trunk depth. Default 30.
+#' @param n_refiner_layers Integer. Refiner depth. Default 2.
+#' @param n_heads Integer. Attention heads. Default 30.
+#' @param norm_eps Numeric. RMSNorm epsilon. Default 1e-5.
+#' @param cap_feat_dim Integer. Caption embedding width. Default 2560.
+#' @param rope_theta Numeric. RoPE base frequency. Default 256.
+#' @param t_scale Numeric. Timestep scale. Default 1000.
+#' @param axes_dims Integer vector. Per-axis rotary dims. Default
+#'   c(32, 48, 48).
+#' @param patch_size Integer. Spatial patch size. Default 2.
+#' @param f_patch_size Integer. Temporal patch size. Default 1.
+#'
+#' @export
+zimage_transformer <- torch::nn_module(
+                                       "zimage_transformer",
+                                       initialize = function(in_channels = 16L, dim = 3840L, n_layers = 30L,
+        n_refiner_layers = 2L, n_heads = 30L, norm_eps = 1e-5,
+        cap_feat_dim = 2560L, rope_theta = 256, t_scale = 1000,
+        axes_dims = c(32L, 48L, 48L), patch_size = 2L, f_patch_size = 1L) {
+    stopifnot(dim %/% n_heads == sum(axes_dims))
+    self$in_channels <- in_channels
+    self$out_channels <- in_channels
+    self$dim <- dim
+    self$rope_theta <- rope_theta
+    self$t_scale <- t_scale
+    self$axes_dims <- axes_dims
+    self$patch_size <- patch_size
+    self$f_patch_size <- f_patch_size
+    patch_key <- paste0(patch_size, "-", f_patch_size)
+    patch_dim <- f_patch_size * patch_size * patch_size * in_channels
+
+    embedders <- list(torch::nn_linear(patch_dim, dim, bias = TRUE))
+    names(embedders) <- patch_key
+    self$all_x_embedder <- torch::nn_module_dict(embedders)
+
+    finals <- list(zimage_final_layer(dim, patch_dim))
+    names(finals) <- patch_key
+    self$all_final_layer <- torch::nn_module_dict(finals)
+
+    self$noise_refiner <- torch::nn_module_list(lapply(
+            seq_len(n_refiner_layers),
+            function(i) zimage_block(dim, n_heads, norm_eps, modulation = TRUE)
+        ))
+    self$context_refiner <- torch::nn_module_list(lapply(
+            seq_len(n_refiner_layers),
+            function(i) zimage_block(dim, n_heads, norm_eps, modulation = FALSE)
+        ))
+    self$layers <- torch::nn_module_list(lapply(seq_len(n_layers),
+            function(i) zimage_block(dim, n_heads, norm_eps, modulation = TRUE)))
+
+    self$t_embedder <- zimage_t_embedder(out_size = min(dim, 256L),
+        mid_size = 1024L)
+    self$cap_embedder <- torch::nn_sequential(
+        ltx23_rms_norm(cap_feat_dim, eps = norm_eps),
+        torch::nn_linear(cap_feat_dim, dim, bias = TRUE)
+    )
+    self$x_pad_token <- torch::nn_parameter(torch::torch_zeros(1L, dim))
+    self$cap_pad_token <- torch::nn_parameter(torch::torch_zeros(1L, dim))
+},
+                                       forward = function(x, t, cap_feats, chunk_size = NULL) {
+    # x: [C, F, H, W] latent; t: [1] in [0, 1]; cap_feats: [L, cap_feat_dim]
+    device <- x$device
+    p <- self$patch_size
+    pf <- self$f_patch_size
+    patch_key <- paste0(p, "-", pf)
+    size <- x$shape[2:4]
+    f_tokens <- size[1] %/% pf
+    h_tokens <- size[2] %/% p
+    w_tokens <- size[3] %/% p
+
+    adaln <- self$t_embedder(t * self$t_scale)$to(dtype = x$dtype) # [1, 256]
+
+    # Position ids and rotary frequencies
+    cap_len <- cap_feats$shape[1]
+    cap_padded <- cap_len + zimage_pad_len(cap_len)
+    cap_freqs <- zimage_pos_embed(
+                                  zimage_cap_pos_ids(cap_padded, device = device),
+                                  axes_dim = self$axes_dims, theta = self$rope_theta
+    )
+    img_freqs <- zimage_pos_embed(
+                                  zimage_img_pos_ids(h_tokens, w_tokens, start0 = cap_padded + 1L,
+            f_tokens = f_tokens, device = device),
+                                  axes_dim = self$axes_dims, theta = self$rope_theta
+    )
+
+    # Image tokens: patchify, embed, pad, refine
+    tokens <- zimage_patchify(x, p, pf)
+    img_len <- tokens$shape[1]
+    x_emb <- self$all_x_embedder[[patch_key]](tokens)$unsqueeze(1L)
+    img_pad <- zimage_pad_len(img_len)
+    if (img_pad > 0L) {
+        pad_rows <- self$x_pad_token$unsqueeze(1L)$expand(c(1L, img_pad, self$dim))
+        x_emb <- torch::torch_cat(list(x_emb, pad_rows), dim = 2L)
+    }
+    for (i in seq_len(length(self$noise_refiner))) {
+        x_emb <- self$noise_refiner[[i]](x_emb, img_freqs, adaln_input = adaln,
+                                       chunk_size = chunk_size)
+    }
+
+    # Caption tokens: embed, pad, refine
+    cap_emb <- self$cap_embedder(cap_feats)$unsqueeze(1L)
+    if (cap_padded > cap_len) {
+        pad_rows <- self$cap_pad_token$unsqueeze(1L)$expand(
+            c(1L, cap_padded - cap_len, self$dim)
+        )
+        cap_emb <- torch::torch_cat(list(cap_emb, pad_rows), dim = 2L)
+    }
+    for (i in seq_len(length(self$context_refiner))) {
+        cap_emb <- self$context_refiner[[i]](cap_emb, cap_freqs,
+                                       chunk_size = chunk_size)
+    }
+
+    # Unified sequence, image first
+    unified <- torch::torch_cat(list(x_emb, cap_emb), dim = 2L)
+    unified_freqs <- list(
+                          torch::torch_cat(list(img_freqs[[1]], cap_freqs[[1]]), dim = 1L),
+                          torch::torch_cat(list(img_freqs[[2]], cap_freqs[[2]]), dim = 1L)
+    )
+    for (i in seq_len(length(self$layers))) {
+        unified <- self$layers[[i]](unified, unified_freqs,
+                                       adaln_input = adaln,
+                                       chunk_size = chunk_size)
+    }
+
+    out <- self$all_final_layer[[patch_key]](unified, adaln)
+    zimage_unpatchify(out$squeeze(1L), size, p, pf, self$out_channels)
+}
+)