diff --git a/CLAUDE.md b/CLAUDE.md index c9603d9..22ebe0f 100644 --- a/CLAUDE.md +++ b/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 diff --git a/DESCRIPTION b/DESCRIPTION index 5f6c71c..5f11d20 100644 --- a/DESCRIPTION +++ b/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")), diff --git a/NAMESPACE b/NAMESPACE index aaeb02f..1d44233 100644 --- a/NAMESPACE +++ b/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) diff --git a/R/checkpoint_flux.R b/R/checkpoint_flux.R index 2012cb4..fc9f551 100644 --- a/R/checkpoint_flux.R +++ b/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" } diff --git a/R/dit_zimage.R b/R/dit_zimage.R new file mode 100644 index 0000000..9776dab --- /dev/null +++ b/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) +} +) diff --git a/R/dit_zimage_modules.R b/R/dit_zimage_modules.R new file mode 100644 index 0000000..64e3c35 --- /dev/null +++ b/R/dit_zimage_modules.R @@ -0,0 +1,157 @@ +#' Z-Image Transformer Block Modules +#' +#' Fresh R port of the Z-Image DiT building blocks from the diffusers +#' reference (Apache-2.0, +#' src/diffusers/models/transformers/transformer_z_image.py). Z-Image is +#' a single-stream DiT: text and image tokens share one sequence and one +#' set of block weights. Each block uses sandwich RMSNorms (a learned +#' norm before AND after both the attention and the feed-forward) and a +#' scale/gate-only modulation — four chunks (scale_msa, gate_msa, +#' scale_mlp, gate_mlp), no shift, gates tanh-squashed, scales 1 + x. +#' The attention is plain joint self-attention, so the FLUX attention +#' module is reused with bias = FALSE and eps = 1e-5. +#' +#' @name dit_zimage_modules +NULL + +# Modulation embedding width (ADALN_EMBED_DIM in the reference); the +# effective width is min(dim, 256). +.zimage_adaln_dim <- 256L + +#' Z-Image feed-forward (SwiGLU with separate gate weights) +#' +#' w2(silu(w1(x)) * w3(x)) with all three linears bias-free. The hidden +#' width is int(dim / 3 * 8). +#' +#' @param dim Integer. Model width. +#' @param hidden_dim Integer. Hidden width. +#' +#' @export +zimage_feed_forward <- torch::nn_module( + "zimage_feed_forward", + initialize = function(dim, hidden_dim) { + self$w1 <- torch::nn_linear(dim, hidden_dim, bias = FALSE) + self$w2 <- torch::nn_linear(hidden_dim, dim, bias = FALSE) + self$w3 <- torch::nn_linear(dim, hidden_dim, bias = FALSE) +}, + forward = function(x) { + self$w2(torch::nnf_silu(self$w1(x)) * self$w3(x)) +} +) + +#' Z-Image transformer block +#' +#' Sandwich-norm residual block shared by the noise refiner, the context +#' refiner and the main trunk. With \code{modulation = TRUE} the block +#' carries an adaLN linear producing (scale_msa, gate_msa, scale_mlp, +#' gate_mlp); the context refiner uses \code{modulation = FALSE} and has +#' no adaLN weights at all. +#' +#' @param dim Integer. Model width. +#' @param n_heads Integer. Attention heads; head dim is dim / n_heads. +#' @param norm_eps Numeric. RMSNorm epsilon. Default 1e-5. +#' @param modulation Logical. Whether the block is timestep-modulated. +#' +#' @export +zimage_block <- torch::nn_module( + "zimage_block", + initialize = function(dim, n_heads, norm_eps = 1e-5, + modulation = TRUE) { + self$attention <- flux_attention(query_dim = dim, heads = n_heads, + dim_head = dim %/% n_heads, eps = 1e-5, + bias = FALSE) + self$feed_forward <- zimage_feed_forward(dim, as.integer(dim / 3 * 8)) + + self$attention_norm1 <- ltx23_rms_norm(dim, eps = norm_eps) + self$ffn_norm1 <- ltx23_rms_norm(dim, eps = norm_eps) + self$attention_norm2 <- ltx23_rms_norm(dim, eps = norm_eps) + self$ffn_norm2 <- ltx23_rms_norm(dim, eps = norm_eps) + + self$modulation <- modulation + if (modulation) { + self$adaLN_modulation <- torch::nn_sequential( + torch::nn_linear(min(dim, .zimage_adaln_dim), 4L * dim, bias = TRUE) + ) + } +}, + forward = function(x, freqs, adaln_input = NULL, chunk_size = NULL) { + if (self$modulation) { + mod <- self$adaLN_modulation(adaln_input)$unsqueeze(2L) + chunks <- mod$chunk(4L, dim = 3L) + scale_msa <- chunks[[1]]$add(1) + gate_msa <- chunks[[2]]$tanh() + scale_mlp <- chunks[[3]]$add(1) + gate_mlp <- chunks[[4]]$tanh() + + attn_out <- self$attention(self$attention_norm1(x) * scale_msa, + image_rotary_emb = freqs, + chunk_size = chunk_size) + x <- x + gate_msa * self$attention_norm2(attn_out) + x + gate_mlp * self$ffn_norm2( + self$feed_forward(self$ffn_norm1(x) * scale_mlp) + ) + } else { + attn_out <- self$attention(self$attention_norm1(x), + image_rotary_emb = freqs, + chunk_size = chunk_size) + x <- x + self$attention_norm2(attn_out) + x + self$ffn_norm2(self$feed_forward(self$ffn_norm1(x))) + } +} +) + +#' Z-Image final layer +#' +#' Parameterless LayerNorm scaled by 1 + adaLN(c) (scale only, no +#' shift), then the token-to-patch projection. +#' +#' @param hidden_size Integer. Model width. +#' @param out_channels Integer. Patch output dim +#' (patch^2 * f_patch * latent channels). +#' +#' @export +zimage_final_layer <- torch::nn_module( + "zimage_final_layer", + initialize = function(hidden_size, out_channels) { + self$norm_final <- torch::nn_layer_norm(hidden_size, eps = 1e-6, + elementwise_affine = FALSE) + self$linear <- torch::nn_linear(hidden_size, out_channels, bias = TRUE) + self$adaLN_modulation <- torch::nn_sequential( + torch::nn_silu(), + torch::nn_linear(min(hidden_size, .zimage_adaln_dim), hidden_size, + bias = TRUE) + ) +}, + forward = function(x, c) { + scale <- self$adaLN_modulation(c)$add(1)$unsqueeze(2L) + self$linear(self$norm_final(x) * scale) +} +) + +#' Z-Image timestep embedder +#' +#' 256-dim cos-first sinusoid (computed in float32) through a +#' Linear-SiLU-Linear MLP. The model feeds t * t_scale with the +#' pipeline's t already in [0, 1]. +#' +#' @param out_size Integer. Output width, min(dim, 256). +#' @param mid_size Integer. Hidden width. The full model uses 1024. +#' @param freq_size Integer. Sinusoid width. Default 256. +#' +#' @export +zimage_t_embedder <- torch::nn_module( + "zimage_t_embedder", + initialize = function(out_size, mid_size = 1024L, freq_size = 256L) { + self$freq_size <- freq_size + self$mlp <- torch::nn_sequential( + torch::nn_linear(freq_size, mid_size, bias = TRUE), + torch::nn_silu(), + torch::nn_linear(mid_size, out_size, bias = TRUE) + ) +}, + forward = function(t) { + t_freq <- ltx23_get_timestep_embedding(t, self$freq_size, + flip_sin_to_cos = TRUE, downscale_freq_shift = 0) + self$mlp(t_freq$to(dtype = self$mlp[[1]]$weight$dtype)) +} +) diff --git a/R/download_zimage.R b/R/download_zimage.R new file mode 100644 index 0000000..08edd4c --- /dev/null +++ b/R/download_zimage.R @@ -0,0 +1,145 @@ +#' Download and Prepare Z-Image-Turbo Weights +#' +#' Downloads Z-Image-Turbo from HuggingFace (Apache-2.0, ungated) and +#' quantizes the 6B transformer to a local fp8 (~6.3 GB) or NF4 +#' (~3.6 GB) artifact. The checkpoint ships the transformer in float32 +#' (24.6 GB), so the one-time quantize saves a lot of disk and load +#' time. +#' +#' @name download_zimage +NULL + +.zimage_repo <- "Tongyi-MAI/Z-Image-Turbo" + +.zimage_transformer_files <- c( + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors.index.json", + sprintf("transformer/diffusion_pytorch_model-%05d-of-00003.safetensors", + 1:3) +) + +.zimage_support_files <- c( + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + "text_encoder/config.json", + "text_encoder/model.safetensors.index.json", + sprintf("text_encoder/model-%05d-of-00003.safetensors", 1:3), + "tokenizer/tokenizer.json", + "tokenizer/tokenizer_config.json", + "scheduler/scheduler_config.json" +) + +#' Download Z-Image-Turbo and build the quantized artifact +#' +#' Skips work already done: a valid quantized manifest short-circuits +#' the transformer download; cached files are not re-fetched. No token +#' is needed (the repo is ungated). The float32 transformer source +#' (~24.6 GB in the HuggingFace cache) may be deleted after +#' quantization. +#' +#' @param quantize Logical. Build the quantized artifact. +#' @param precision "fp8" (~6.3 GB, GPU-resident; near-bf16 quality) or +#' "nf4" (~3.6 GB). +#' @param output_dir Directory for the quantized artifact. +#' @param text_encoders Logical. Also fetch the Qwen3-4B text encoder, +#' tokenizer, VAE, and scheduler config (~8.2 GB). +#' @param verbose Logical. +#' +#' @return Invisibly, a list with \code{transformer_dir}, +#' \code{artifact_dir}, and \code{support} (named file paths). +#' +#' @export +download_zimage_turbo <- function(quantize = TRUE, + precision = c("fp8", "nf4"), + output_dir = NULL, text_encoders = TRUE, + verbose = TRUE) { + precision <- match.arg(precision) + if (is.null(output_dir)) { + output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + paste0("zimage-turbo-", precision)) + } + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to download model weights.") + } + result <- list(transformer_dir = NULL, artifact_dir = output_dir, + support = character(0)) + + manifest_path <- file.path(output_dir, "manifest.json") + have_artifact <- file.exists(manifest_path) && { + m <- jsonlite::fromJSON(manifest_path) + all(file.exists(file.path(output_dir, m$shards))) + } + + if (!have_artifact || !quantize) { + cached <- tryCatch( + hfhub::hub_download(.zimage_repo, .zimage_transformer_files[[3]], + local_files_only = TRUE), + error = function(e) NULL + ) + if (is.null(cached) && !have_artifact) { + free <- .ltx23_disk_free_gb(path.expand("~")) + if (!is.na(free) && free < 35) { + warning(sprintf( + "Only %.0f GB free; the download + %s artifact need ~35 GB.", + free, precision + )) + } + ok <- .ltx23_consent(paste0( + "Z-Image-Turbo: the 24.6 GB float32 transformer plus a local ", + precision, " artifact (Apache-2.0, ungated)" + )) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading the Z-Image-Turbo transformer (24.6 GB)...") + } + } + paths <- vapply(.zimage_transformer_files, function(f) { + hfhub::hub_download(.zimage_repo, f) + }, character(1)) + result$transformer_dir <- dirname(paths[[1]]) + + if (quantize && !have_artifact) { + if (verbose) { + message("Quantizing transformer linears to ", precision, + " (one-time)...") + } + flux_quantize(result$transformer_dir, output_dir, + format = precision, verbose = verbose) + if (verbose) { + message( + toupper(precision), " artifact ready: ", output_dir, "\n", + "The 24.6 GB float32 source in the HuggingFace cache ", + "may be deleted if you do not need it." + ) + } + } + } else if (verbose) { + message(toupper(precision), " artifact already present: ", output_dir) + } + + if (text_encoders) { + have_te <- !is.null(tryCatch( + hfhub::hub_download(.zimage_repo, .zimage_support_files[[5]], + local_files_only = TRUE), + error = function(e) NULL + )) + if (!have_te) { + ok <- .ltx23_consent( + "the Qwen3-4B text encoder, tokenizer, and VAE (~8.2 GB)" + ) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading text encoder + VAE...") + } + } + result$support <- vapply(.zimage_support_files, function(f) { + hfhub::hub_download(.zimage_repo, f) + }, character(1)) + } + + invisible(result) +} diff --git a/R/memory_flux.R b/R/memory_flux.R index 14f8b97..3e7c7ec 100644 --- a/R/memory_flux.R +++ b/R/memory_flux.R @@ -8,6 +8,44 @@ #' @name memory_flux NULL +# Allocator gc gates for the resident-weight image pipelines +# (backend:native). The allocated/reserved ratio gate is chronically +# over its 0.95 threshold under a lazily-reserving caching allocator +# (measured 0.957 at steady state), so the R gc callback fired on +# nearly every allocation - ~2,200 gcs x 62 ms = 89% of wall time on a +# Z-Image generation. Disabling the ratio gate alone OOMs (dead tensor +# handles accumulate until finalizers run), so the absolute +# allocated_rate gate takes over garbage-accumulation duty at 0.65 of +# total VRAM, above the ~8.3 GB live phase peak with headroom under the +# card limit. Measured (RTX 5060 Ti 16 GB): Z-Image 512^2 142 -> 12 s, +# 1024^2 143 -> 25 s; klein 1024^2 48 -> 13 s, no OOM, reserved peak +# 12.9 GB. ltx23_tune_gc's reserved_rate stays as the fragmentation +# net; the LTX video pipelines keep their own separately measured +# gates. +.flux_gc_gates <- function(footprint_gb = 12) { + if (is.null(getOption("torch.cuda_allocator_allocated_reserved_rate"))) { + options(torch.cuda_allocator_allocated_reserved_rate = 1.0) + } + if (is.null(getOption("torch.cuda_allocator_allocated_rate"))) { + options(torch.cuda_allocator_allocated_rate = 0.65) + } + ltx23_tune_gc(footprint_gb = footprint_gb) + # start_torch() reads the three gate options ONCE at torch init and + # pushes them into the C++ allocator; option changes after that are + # inert. Torch is long started by the time a loader runs, so push + # the current values into the live allocator directly. + push <- get0("cpp_set_cuda_allocator_allocator_thresholds", + envir = asNamespace("torch")) + if (is.function(push)) { + try(push( + getOption("torch.cuda_allocator_reserved_rate", 0.2), + getOption("torch.cuda_allocator_allocated_rate", 0.8), + getOption("torch.cuda_allocator_allocated_reserved_rate", 0.8) + ), silent = TRUE) + } + invisible(NULL) +} + #' Resolve a FLUX memory profile #' #' @param vram_gb Numeric or NULL. Available VRAM; auto-detected when diff --git a/R/models2devices.R b/R/models2devices.R index 7b5f3f9..e0f6bb8 100644 --- a/R/models2devices.R +++ b/R/models2devices.R @@ -59,7 +59,8 @@ get_required_components <- function(model_name) { "sdxl" = c("unet", "decoder", "text_encoder", "text_encoder2", "encoder"), "flux1" = c("transformer", "decoder", "text_encoder", "text_encoder2"), - "flux2" = c("transformer", "decoder", "text_encoder") + "flux2" = c("transformer", "decoder", "text_encoder"), + "zimage" = c("transformer", "decoder", "text_encoder") # "sd3" = c("transformer", "decoder", "text_encoder", "text_encoder2", "text_encoder3", "encoder"), # "cascade" = c("prior", "decoder", "text_encoder", "vqgan") ) diff --git a/R/quantize_flux.R b/R/quantize_flux.R index 6f4376e..ce1f3bb 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -71,6 +71,28 @@ NULL args } +.zimage_transformer_args <- function(config) { + if (is.null(config)) { + return(list()) + } + args <- list(in_channels = config$in_channels, dim = config$dim, + n_layers = config$n_layers, + n_refiner_layers = config$n_refiner_layers, + n_heads = config$n_heads, cap_feat_dim = config$cap_feat_dim, + axes_dims = config$axes_dims, + patch_size = config$all_patch_size[[1]], + f_patch_size = config$all_f_patch_size[[1]]) + args <- Filter(function(x) !is.null(x) && length(x) > 0L, args) + args <- lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x) + for (field in c("norm_eps", "rope_theta", "t_scale")) { + v <- config[[field]] + if (!is.null(v) && length(v) == 1L) { + args[[field]] <- as.numeric(v) + } + } + args +} + # Move plain-field fp8 weights (and their scales) onto a device; used # for resident fp8 where the whole quantized model fits on the GPU .flux_fp8_to_device <- function(module, device) { @@ -87,9 +109,15 @@ NULL # Family-specific hooks for quantization and loading .flux_family_hooks <- function(config) { - if (.flux_family(config) == "flux2") { + family <- .flux_family(config) + if (family == "flux2") { list(model_fn = flux2_transformer, args_fn = .flux2_transformer_args, is_quant_key = flux2_is_quant_key, shard_prefix = "flux2-klein") + } else if (family == "zimage") { + list(model_fn = zimage_transformer, + args_fn = .zimage_transformer_args, + is_quant_key = zimage_is_quant_key, + shard_prefix = "zimage-turbo") } else { list(model_fn = flux_transformer, args_fn = .flux_transformer_args, @@ -192,7 +220,10 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, }) n_cast <- n_cast + 1L } else { - shard[[key]] <- tensor + # Residents load into the compute dtype anyway; storing bf16 + # keeps fp32-shipped checkpoints (Z-Image) from doubling the + # artifact. Identity for bf16 sources. + shard[[key]] <- tensor$to(dtype = torch::torch_bfloat16()) shard_size <- shard_size + prod(tensor$shape) * 2 } rm(tensor) diff --git a/R/rope_zimage.R b/R/rope_zimage.R new file mode 100644 index 0000000..c8c553c --- /dev/null +++ b/R/rope_zimage.R @@ -0,0 +1,195 @@ +#' Z-Image Rotary Positional Embeddings and Patchify Helpers +#' +#' Fresh R port of the Z-Image position scheme from the diffusers +#' reference (Apache-2.0, +#' src/diffusers/models/transformers/transformer_z_image.py RopeEmbedder, +#' create_coordinate_grid, _patchify_image, _pad_with_ids, unpatchify). +#' Z-Image uses 3-axis interleaved RoPE with theta 256; frequencies are +#' built in float64 but the angles are cast to float32 before cos/sin +#' (torch.polar on a .float() tensor), which differs measurably from the +#' FLUX convention at large positions. Every sub-sequence is padded to a +#' multiple of 32 (SEQ_MULTI_OF); caption positions are a 1-based ramp on +#' axis 1 built over the padded length, image positions sit on axes 2/3 +#' with axis 1 offset just past the caption. +#' +#' @name rope_zimage +NULL + +# Sub-sequences are padded to a multiple of this many tokens. +.zimage_seq_multi_of <- 32L + +#' Padding length to the next multiple of 32 +#' +#' @param n Integer token count. +#' @return Integer pad length in [0, 31]. +#' @keywords internal +zimage_pad_len <- function(n) { + (-n) %% .zimage_seq_multi_of +} + +#' Build Z-Image caption position ids +#' +#' Caption tokens ramp 1..cap_padded_len on the first axis (axes 2 and 3 +#' zero). The reference builds the grid over the already-padded length, +#' so pad tokens continue the ramp rather than sitting at the origin +#' (the (0,0,0) pad ids it also emits are truncated away in +#' _prepare_sequence and never reach RoPE). +#' +#' @param cap_padded_len Integer. Caption length after padding to a +#' multiple of 32. +#' @param device Device for the resulting tensor. +#' +#' @return Float tensor of shape [cap_padded_len, 3]. +#' +#' @export +zimage_cap_pos_ids <- function(cap_padded_len, device = "cpu") { + f32 <- torch::torch_float32() + ids <- torch::torch_zeros(cap_padded_len, 3L, dtype = f32, device = device) + ids[, 1] <- torch::torch_arange(start = 1, end = cap_padded_len, + dtype = f32, device = device) + ids +} + +#' Build Z-Image latent image position ids +#' +#' Image tokens use axis 1 for the frame index offset past the caption +#' (start0 = cap_padded_len + 1), axis 2 for the token row and axis 3 for +#' the token column. Trailing pad tokens (token count not a multiple of +#' 32) sit at (0, 0, 0). Reference: patchify_and_embed / _pad_with_ids. +#' +#' @param h_tokens Integer. Token grid height (latent height / patch). +#' @param w_tokens Integer. Token grid width (latent width / patch). +#' @param start0 Integer. First-axis start, cap_padded_len + 1. +#' @param f_tokens Integer. Token grid frames; 1 for txt2img. +#' @param device Device for the resulting tensor. +#' +#' @return Float tensor of shape [padded token count, 3]. +#' +#' @export +zimage_img_pos_ids <- function(h_tokens, w_tokens, start0, f_tokens = 1L, + device = "cpu") { + f32 <- torch::torch_float32() + ids <- torch::torch_zeros(f_tokens, h_tokens, w_tokens, 3L, + dtype = f32, device = device) + frames <- torch::torch_arange(start = start0, + end = start0 + f_tokens - 1, + dtype = f32, device = device) + rows <- torch::torch_arange(start = 0, end = h_tokens - 1, dtype = f32, + device = device) + cols <- torch::torch_arange(start = 0, end = w_tokens - 1, dtype = f32, + device = device) + ids[,,, 1] <- ids[,,, 1] + frames$reshape(c(-1L, 1L, 1L)) + ids[,,, 2] <- ids[,,, 2] + rows$reshape(c(1L, -1L, 1L)) + ids[,,, 3] <- ids[,,, 3] + cols$reshape(c(1L, 1L, -1L)) + ids <- ids$reshape(c(f_tokens * h_tokens * w_tokens, 3L)) + + pad <- zimage_pad_len(ids$shape[1]) + if (pad > 0L) { + ids <- torch::torch_cat(list( + ids, + torch::torch_zeros(pad, 3L, dtype = f32, device = device) + )) + } + ids +} + +#' Compute Z-Image rotary frequencies from position ids +#' +#' Per-axis 1D rotary frequencies in the interleaved-real convention. +#' Frequencies and angles are built in float64, then the angles are cast +#' to float32 before cos/sin — matching the reference torch.polar call on +#' a .float() tensor. Output format matches \code{flux_pos_embed} so +#' \code{flux_apply_rotary_emb} applies unchanged. +#' +#' @param ids Tensor of shape [S, 3] from \code{zimage_cap_pos_ids} / +#' \code{zimage_img_pos_ids}. +#' @param axes_dim Integer vector of per-axis rotary dims; must sum to +#' the attention head dim. Z-Image uses c(32, 48, 48). +#' @param theta Numeric. RoPE base frequency. Z-Image uses 256. +#' +#' @return List of two tensors (cos, sin), each [S, sum(axes_dim)], +#' float32, on the device of \code{ids}. +#' +#' @export +zimage_pos_embed <- function(ids, axes_dim = c(32L, 48L, 48L), theta = 256) { + n_axes <- ids$shape[2] + device <- ids$device + f64 <- torch::torch_float64() + f32 <- torch::torch_float32() + pos <- ids$to(dtype = f64)$cpu() + + cos_out <- vector("list", n_axes) + sin_out <- vector("list", n_axes) + for (i in seq_len(n_axes)) { + d <- axes_dim[i] + exponents <- torch::torch_arange(start = 0, end = d - 2, step = 2, + dtype = f64) + freqs <- 1.0 / torch::torch_pow(theta, exponents / d) + # Angle in float64, cast to float32 BEFORE cos/sin (reference: + # torch.outer(...).float() then torch.polar) + angles <- (pos[, i]$unsqueeze(2L) * freqs$unsqueeze(1L))$to(dtype = f32) + cos_out[[i]] <- angles$cos()$repeat_interleave(2L, dim = 2L) + sin_out[[i]] <- angles$sin()$repeat_interleave(2L, dim = 2L) + } + + list( + torch::torch_cat(cos_out, dim = -1L)$to(device = device), + torch::torch_cat(sin_out, dim = -1L)$to(device = device) + ) +} + +#' Patchify a latent image to Z-Image tokens +#' +#' (C, F, H, W) -> [F/pF * H/p * W/p, pF * p * p * C], matching +#' _patchify_image. No padding is applied here. +#' +#' @param image Tensor of shape [C, F, H, W]. +#' @param patch_size Integer spatial patch size. Default 2. +#' @param f_patch_size Integer temporal patch size. Default 1. +#' +#' @return Tensor of shape [num_tokens, patch_dim]. +#' +#' @export +zimage_patchify <- function(image, patch_size = 2L, f_patch_size = 1L) { + p <- patch_size + pf <- f_patch_size + shape <- image$shape + c_in <- shape[1] + f_tokens <- shape[2] %/% pf + h_tokens <- shape[3] %/% p + w_tokens <- shape[4] %/% p + image <- image$view(c(c_in, f_tokens, pf, h_tokens, p, w_tokens, p)) + # Python permute(1, 3, 5, 2, 4, 6, 0) -> R 1-indexed + image <- image$permute(c(2L, 4L, 6L, 3L, 5L, 7L, 1L)) + image$reshape(c(f_tokens * h_tokens * w_tokens, pf * p * p * c_in)) +} + +#' Unpatchify Z-Image tokens back to a latent image +#' +#' Takes the first F/pF * H/p * W/p tokens (the image span of the +#' unified sequence) and reassembles [C, F, H, W], matching unpatchify. +#' +#' @param tokens Tensor of shape [S, pF * p * p * C] with the image +#' tokens first. +#' @param size Integer vector c(F, H, W) of the target latent size. +#' @param patch_size Integer spatial patch size. Default 2. +#' @param f_patch_size Integer temporal patch size. Default 1. +#' @param out_channels Integer number of latent channels. Default 16. +#' +#' @return Tensor of shape [C, F, H, W]. +#' +#' @export +zimage_unpatchify <- function(tokens, size, patch_size = 2L, + f_patch_size = 1L, out_channels = 16L) { + p <- patch_size + pf <- f_patch_size + f_tokens <- size[1] %/% pf + h_tokens <- size[2] %/% p + w_tokens <- size[3] %/% p + ori_len <- f_tokens * h_tokens * w_tokens + x <- tokens[1:ori_len,] + x <- x$view(c(f_tokens, h_tokens, w_tokens, pf, p, p, out_channels)) + # Python permute(6, 0, 3, 1, 4, 2, 5) -> R 1-indexed + x <- x$permute(c(7L, 1L, 4L, 2L, 5L, 3L, 6L)) + x$reshape(c(out_channels, size[1], size[2], size[3])) +} diff --git a/R/tokenizer_qwen.R b/R/tokenizer_qwen.R index 0ab7a9d..6822a08 100644 --- a/R/tokenizer_qwen.R +++ b/R/tokenizer_qwen.R @@ -201,17 +201,21 @@ print.qwen_tokenizer <- function(x, ...) { #' Encode prompts with the Qwen tokenizer #' -#' With \code{chat_template = TRUE} (the FLUX.2 klein pipeline behavior) -#' each prompt is wrapped as a single user turn with the generation -#' prompt and a disabled thinking block, matching -#' \code{apply_chat_template(..., add_generation_prompt = TRUE, -#' enable_thinking = FALSE)}. Right-pads with \code{<|endoftext|>}. +#' With \code{chat_template = TRUE} each prompt is wrapped as a single +#' user turn with the generation prompt, matching +#' \code{apply_chat_template(..., add_generation_prompt = TRUE)}. With +#' \code{enable_thinking = FALSE} (the FLUX.2 klein pipeline behavior) +#' the template closes with an empty thinking block; with +#' \code{enable_thinking = TRUE} (the Z-Image pipeline behavior) it ends +#' at the assistant turn. Right-pads with \code{<|endoftext|>}. #' #' @param tokenizer A \code{\link{qwen_bpe_tokenizer}}. #' @param texts Character vector of prompts. #' @param max_length Integer. Fixed sequence length (klein: 512). NULL #' for no truncation/padding. #' @param chat_template Logical. Wrap in the Qwen3 chat template. +#' @param enable_thinking Logical. Leave the model's thinking enabled +#' (no empty think block). Default FALSE. #' #' @return List with \code{input_ids} and \code{attention_mask} integer #' matrices [length(texts), max_length] (ragged lists when @@ -219,13 +223,14 @@ print.qwen_tokenizer <- function(x, ...) { #' #' @export encode_qwen <- function(tokenizer, texts, max_length = 512L, - chat_template = TRUE) { + chat_template = TRUE, enable_thinking = FALSE) { stopifnot(inherits(tokenizer, "qwen_tokenizer")) encode_one <- function(text) { if (chat_template) { text <- paste0("<|im_start|>user\n", text, "<|im_end|>\n", - "<|im_start|>assistant\n\n\n\n\n") + "<|im_start|>assistant\n", + if (!enable_thinking) "\n\n\n\n") } ids <- integer(0) for (chunk in .qwen_split_added(text, tokenizer)) { diff --git a/R/txt2img.R b/R/txt2img.R index 2d1742b..3384fc1 100644 --- a/R/txt2img.R +++ b/R/txt2img.R @@ -11,7 +11,8 @@ #' \dontrun{ #' img <- txt2img("a cat wearing sunglasses in space", device = "cuda") #' } -txt2img <- function(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2"), +txt2img <- function(prompt, + model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage"), ...) { switch(model_name, # "sd15" = txt2img_sd15(prompt, ...), @@ -19,6 +20,7 @@ txt2img <- function(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2"), "sdxl" = txt2img_sdxl(prompt, ...), "flux1" = txt2img_flux(prompt, ...), "flux2" = txt2img_flux2(prompt, ...), + "zimage" = txt2img_zimage(prompt, ...), # "sd3" = txt2img_sd3(prompt, ...), stop("Unsupported model: ", model_name) ) diff --git a/R/txt2img_flux.R b/R/txt2img_flux.R index 1bf9282..857489f 100644 --- a/R/txt2img_flux.R +++ b/R/txt2img_flux.R @@ -131,7 +131,7 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", } else { footprint <- 4 } - ltx23_tune_gc(footprint_gb = footprint) + .flux_gc_gates(footprint_gb = footprint) } if (phase_offload) { diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R index 30241f1..7f42705 100644 --- a/R/txt2img_flux2.R +++ b/R/txt2img_flux2.R @@ -76,12 +76,12 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "backend:native") } if (device == "cuda") { - # Sized to the LARGEST phase (the 8 GB Qwen3 encode), not the - # transformer: a low footprint puts the allocator's R-gc - # callback threshold under the working set, and every callback - # walks the ~300k-object tokenizer heap (measured: 13-20 s - # forwards at footprint 6 vs sub-second at 12) - ltx23_tune_gc(footprint_gb = 12) + # Footprint sized to the LARGEST phase (the 8 GB Qwen3 encode), + # not the transformer: a low footprint puts the allocator's + # R-gc callback threshold under the working set, and every + # callback walks the ~300k-object tokenizer heap (measured: + # 13-20 s forwards at footprint 6 vs sub-second at 12) + .flux_gc_gates(footprint_gb = 12) } if (phase_offload) { diff --git a/R/txt2img_zimage.R b/R/txt2img_zimage.R new file mode 100644 index 0000000..ea51105 --- /dev/null +++ b/R/txt2img_zimage.R @@ -0,0 +1,376 @@ +#' Z-Image-Turbo Text-to-Image Pipeline +#' +#' Z-Image-Turbo text-to-image, ported from the diffusers reference +#' (Apache-2.0, src/diffusers/pipelines/z_image/pipeline_z_image.py). +#' Turbo is guidance-distilled: 8 steps, no classifier-free guidance. +#' The FlowMatch schedule uses the checkpoint's 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. The VAE is the FLUX.1 16-channel autoencoder +#' verbatim. +#' +#' @name txt2img_zimage +NULL + +# Resolve a Z-Image support file from the HuggingFace cache +.zimage_cached <- function(file) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to locate model files.") + } + tryCatch( + hfhub::hub_download(.zimage_repo, file, local_files_only = TRUE), + error = function(e) { + stop("Missing ", file, " in the HuggingFace cache; ", + "run download_zimage_turbo() first.", call. = FALSE) + } + ) +} + +#' Load the Z-Image-Turbo pipeline +#' +#' Loads the quantized transformer artifact plus the 16-channel VAE +#' decoder, Qwen3-4B text encoder, and tokenizer from the HuggingFace +#' cache populated by \code{\link{download_zimage_turbo}}. With fp8 +#' precision the ~6.3 GB transformer rides to the GPU per phase. +#' +#' @param model_dir Quantized artifact directory (default: the +#' \code{download_zimage_turbo} location for \code{precision}), or a +#' raw diffusers transformer directory. +#' @param device Character. Compute device. +#' @param precision "fp8" (default) or "nf4". +#' @param text_device Device for the Qwen3 encoder (default: +#' \code{device}; it encodes in its own phase and offloads). +#' @param attn_chunk Integer or NULL. Attention query-chunk override. +#' @param phase_offload Logical. One GPU tenant per phase. +#' @param verbose Logical. +#' +#' @return A \code{zimage_pipeline} list. +#' +#' @export +zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", + precision = c("fp8", "nf4"), + text_device = NULL, attn_chunk = NULL, + phase_offload = TRUE, verbose = TRUE) { + precision <- match.arg(precision) + if (is.null(text_device)) { + if (device == "cuda") { + text_device <- "cuda" + } else { + text_device <- "cpu" + } + } + if (is.null(model_dir)) { + model_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + paste0("zimage-turbo-", precision)) + } + + ckpt <- if (file.exists(file.path(model_dir, "manifest.json"))) { + flux_open_quantized(model_dir) + } else { + flux_open_checkpoint(model_dir) + } + + if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { + # Stable resident footprint: the native backend avoids + # expandable_segments' page-unmap cost on per-step churn + Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = "backend:native") + } + if (device == "cuda") { + # Footprint sized to the largest phase (the 8 GB Qwen3 encode), + # same reasoning as the FLUX.2 klein pipeline + .flux_gc_gates(footprint_gb = 12) + } + + if (phase_offload) { + component_device <- "cpu" + } else { + component_device <- device + } + + pipe <- list( + format = ckpt$format %||% "full", + device = device, + text_device = text_device, + phase_offload = phase_offload, + attn_chunk = if (is.null(attn_chunk)) NULL else as.integer(attn_chunk), + config = ckpt$config + ) + + if (verbose) { + message("Loading transformer (", pipe$format, ")...") + } + pipe$transformer <- flux_load_transformer( + ckpt, device = component_device, + dtype = if (device == "cpu") "float32" else "bfloat16", + pin = FALSE, + fp8_resident = FALSE, + verbose = verbose + ) + # Resident fp8 happens at onload time (the weights ride to the GPU + # with the phase and back off after) + pipe$fp8_resident <- identical(pipe$format, "fp8") && device == "cuda" + + if (verbose) { + message("Loading VAE decoder...") + } + vae_config <- jsonlite::fromJSON(.zimage_cached("vae/config.json")) + pipe$vae_scaling_factor <- vae_config$scaling_factor %||% 0.3611 + pipe$vae_shift_factor <- vae_config$shift_factor %||% 0.1159 + dec <- vae_decoder_native( + latent_channels = as.integer(vae_config$latent_channels %||% 16L) + ) + load_decoder_safetensors( + dec, .zimage_cached("vae/diffusion_pytorch_model.safetensors"), + verbose = verbose + ) + dec$to(device = component_device) + dec$eval() + pipe$decoder <- dec + + sched_config <- jsonlite::fromJSON(.zimage_cached("scheduler/scheduler_config.json")) + pipe$sched_shift <- sched_config$shift %||% 3.0 + + if (verbose) { + message("Loading Qwen3 text encoder...") + } + te_dir <- dirname(.zimage_cached("text_encoder/config.json")) + te_config <- jsonlite::fromJSON(file.path(te_dir, "config.json")) + # hidden_states[-2]: the state after num_hidden_layers - 1 layers + pipe$te_penult_layer <- as.integer(te_config$num_hidden_layers %||% 36L) - 1L + pipe$text_encoder <- load_qwen3_text_encoder( + te_dir, device = if (phase_offload) "cpu" else text_device, + dtype = if (text_device == "cpu") "float32" else "bfloat16", + verbose = verbose + ) + pipe$tokenizer <- qwen_bpe_tokenizer(.zimage_cached("tokenizer/tokenizer.json")) + + structure(pipe, class = "zimage_pipeline") +} + +# Qwen3 prompt encoding, Z-Image style: thinking-enabled chat template, +# penultimate hidden state, valid tokens only (right padding -> first n) +.zimage_encode_prompt <- function(prompt, model, tokenizer, penult_layer, + max_sequence_length = 512L, device = NULL) { + enc <- encode_qwen(tokenizer, prompt, max_length = max_sequence_length, + chat_template = TRUE, enable_thinking = TRUE) + device <- device %||% model$model$embed_tokens$weight$device + long <- torch::torch_long() + ids <- torch::torch_tensor(enc$input_ids + 1L, dtype = long, + device = device) + mask <- torch::torch_tensor(enc$attention_mask, dtype = long, + device = device) + + states <- torch::with_no_grad(model(ids, attention_mask = mask, + out_layers = penult_layer)) + n_real <- sum(enc$attention_mask[1,]) + states[[1]][1, 1:n_real,] # [L, hidden] +} + +# Flow-matching Euler loop; Turbo is CFG-free (one forward per step). +# The model sees the reversed normalized timestep and predicts the +# negated velocity. +.zimage_denoise <- function(transformer, latents, schedule, cap_feats, + compute_dtype, chunk_size = NULL, verbose = TRUE) { + timesteps <- as.numeric(schedule$timesteps$cpu()) + n <- length(timesteps) + pb <- if (verbose) { + utils::txtProgressBar(min = 0, max = n, style = 3) + } else { + NULL + } + f32 <- torch::torch_float32() + + torch::with_no_grad({ + for (i in seq_len(n)) { + t <- timesteps[[i]] + t_model <- torch::torch_tensor((1000 - t) / 1000, dtype = f32, + device = latents$device)$reshape(1L) + + # [1, 16, H8, W8] -> [16, 1, H8, W8] (frame axis) + x_in <- latents$squeeze(1L)$unsqueeze(2L)$to(dtype = compute_dtype) + out <- transformer(x_in, t_model, cap_feats, + chunk_size = chunk_size) + noise_pred <- out$squeeze(2L)$unsqueeze(1L)$to(dtype = f32)$neg() + + step <- flowmatch_scheduler_step(noise_pred, t, latents, schedule) + latents <- step$prev_sample + schedule <- step$schedule + rm(out, noise_pred, step) + if (!is.null(pb)) { + utils::setTxtProgressBar(pb, i) + } + } + }) + if (!is.null(pb)) { + close(pb) + } + latents +} + +#' Generate an image with Z-Image-Turbo +#' +#' Guidance-distilled text-to-image (8 steps, no CFG): Qwen3-4B prompt +#' encoding (thinking-enabled chat template, penultimate hidden state), +#' FlowMatch denoising with the reversed-timestep convention, and +#' 16-channel VAE decode. Strong at legible text rendering, English and +#' Chinese both. +#' +#' @param prompt Character. The prompt. +#' @param pipeline A \code{zimage_pipeline} from +#' \code{\link{zimage_load_pipeline}}; NULL loads one (passing +#' \code{...} through). +#' @param width,height Integers, divisible by 16. +#' @param num_inference_steps Integer. Denoising steps (Turbo: 8). +#' @param max_sequence_length Integer. Qwen3 token length (512). +#' @param seed Integer or NULL. Latents are drawn on the CPU, so a seed +#' matches a Python diffusers run with a CPU generator. +#' @param prompt_embeds Optional precomputed [L, 2560] caption +#' embeddings (valid tokens only). +#' @param save_file Logical. Write a PNG. +#' @param filename Output path (default derived from the prompt). +#' @param verbose Logical. +#' @param ... Passed to \code{\link{zimage_load_pipeline}} when +#' \code{pipeline} is NULL. +#' +#' @return Invisibly, \code{list(image, metadata)} where \code{image} is +#' an [H, W, 3] array in [0, 1]. +#' +#' @export +txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, + height = 1024L, num_inference_steps = 8L, + max_sequence_length = 512L, seed = NULL, + prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) { + if (is.null(pipeline)) { + pipeline <- zimage_load_pipeline(..., verbose = verbose) + } + device <- pipeline$device + width <- as.integer(width) + height <- as.integer(height) + if (width %% 16L != 0L || height %% 16L != 0L) { + stop("width and height must be divisible by 16") + } + + f32 <- torch::torch_float32() + compute_dtype <- if (device == "cpu") { + f32 + } else { + torch::torch_bfloat16() + } + + phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" + onload <- function(module) { + if (phase_offload) { + module$to(device = device) + } + module + } + offload <- function(module) { + if (phase_offload) { + module$to(device = "cpu") + clear_vram() + } + invisible(module) + } + + t0 <- Sys.time() + + # --- Phase 1: text encoding -------------------------------------------------- + if (is.null(prompt_embeds)) { + if (verbose) { + message("Encoding prompt (Qwen3)...") + } + onload(pipeline$text_encoder) + te_device <- pipeline$text_encoder$model$embed_tokens$weight$device + prompt_embeds <- .zimage_encode_prompt(prompt, pipeline$text_encoder, + pipeline$tokenizer, + penult_layer = pipeline$te_penult_layer %||% 35L, + max_sequence_length = max_sequence_length, device = te_device) + offload(pipeline$text_encoder) + } + prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) + + # --- Phase 2: latents and schedule ---------------------------------------------- + h8 <- height %/% 8L + w8 <- width %/% 8L + if (!is.null(seed)) { + torch::torch_manual_seed(seed) + } + latents <- torch::torch_randn(c(1L, 16L, h8, w8), dtype = f32)$ + to(device = device) + + n_steps <- as.integer(num_inference_steps) + sched <- flowmatch_scheduler_create( + shift = pipeline$sched_shift %||% 3.0, + use_dynamic_shifting = FALSE + ) + sched <- flowmatch_set_timesteps( + sched, n_steps, + sigmas = seq(1, 1 / n_steps, length.out = n_steps) + ) + + # --- Phase 3: denoise ------------------------------------------------------------ + transformer <- onload(pipeline$transformer) + if (isTRUE(pipeline$fp8_resident)) { + .flux_fp8_to_device(transformer, device) + } + if (verbose) { + message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, + height)) + } + latents <- .zimage_denoise( + transformer, latents, sched, prompt_embeds, + compute_dtype, chunk_size = pipeline$attn_chunk, + verbose = verbose + ) + if (isTRUE(pipeline$fp8_resident) && phase_offload) { + .flux_fp8_to_device(pipeline$transformer, "cpu") + } + offload(pipeline$transformer) + ltx23_release_dequant_buffers() + + # --- Phase 4: decode --------------------------------------------------------------- + if (verbose) { + message("Decoding...") + } + latents <- latents$div(pipeline$vae_scaling_factor %||% 0.3611)$ + add(pipeline$vae_shift_factor %||% 0.1159) + + decoder <- pipeline$decoder + if (phase_offload) { + decoder$to(device = device, dtype = compute_dtype) + } + torch::with_no_grad({ + dec_param <- decoder$conv_in$weight + img <- decoder(latents$to(device = dec_param$device, + dtype = dec_param$dtype)) + img <- img$to(dtype = f32)$cpu() + }) + offload(decoder) + + img <- img$squeeze(1)$permute(c(2L, 3L, 1L)) + img <- img$add(1)$div(2)$clamp(0, 1) + img_array <- as.array(img) + + gen_seconds <- as.numeric(difftime(Sys.time(), t0, units = "secs")) + if (verbose) { + message(sprintf("Generated in %.1f s", gen_seconds)) + } + + if (save_file) { + if (is.null(filename)) { + filename <- filename_from_prompt(prompt) + } + save_image(img_array, filename) + if (verbose) { + message("Saved to ", filename) + } + } + + metadata <- list( + prompt = prompt, width = width, height = height, + steps = n_steps, seed = seed, model = "zimage-turbo", + precision = pipeline$format, seconds = gen_seconds + ) + invisible(list(image = img_array, metadata = metadata)) +} diff --git a/README.md b/README.md index 7d96617..77bef08 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ targets::install_github("cornball-ai/diffuseR") ## Features -- **Text-to-Image Generation**: Stable Diffusion 2.1, SDXL, FLUX.1-schnell, and FLUX.2 Klein (fully native R torch implementations) +- **Text-to-Image Generation**: Stable Diffusion 2.1, SDXL, FLUX.1-schnell, FLUX.2 Klein, and Z-Image-Turbo (fully native R torch implementations) - **Text-to-Video Generation**: LTX-2.3 with synchronized audio - **Image-to-Image Generation**: Modify existing images based on text prompts (SD 2.1 / SDXL) - **GPU-poor support**: NF4 and fp8 quantization run the 12B FLUX.1 and 22B LTX-2.3 transformers on a 16GB card @@ -135,13 +135,13 @@ torch::cuda_empty_cache() ![](man/figures/gambling_cat.png) -### FLUX +### FLUX and Z-Image -FLUX.1-schnell (12B) and FLUX.2 Klein (4B) are step-distilled models: -4 denoising steps, no guidance. Both are quantized locally once at -download time and fit comfortably on a 16GB GPU (measured on an RTX -5060 Ti: FLUX.1 1024x1024 in ~2 min at 8.7GB peak; FLUX.2 Klein in -~48s at 8.2GB peak). +FLUX.1-schnell (12B), FLUX.2 Klein (4B), and Z-Image-Turbo (6B) are +step-distilled models: 4-8 denoising steps, no guidance. All are +quantized locally once at download time and fit comfortably on a 16GB +GPU (measured 1024x1024 on an RTX 5060 Ti: FLUX.1 ~55s at 9.6GB peak; +FLUX.2 Klein ~13s at 12.5GB; Z-Image-Turbo ~24s at 13.1GB). ```r library(diffuseR) @@ -159,6 +159,12 @@ download_flux2_klein() txt2img_flux2("a red fox sitting in a snowy forest, digital art", seed = 42) +# Z-Image-Turbo: ungated, strong at legible text in images (EN + CN). +# ~33GB download, one-time fp8 quantize to a 5.9GB artifact. +download_zimage_turbo() +txt2img_zimage(paste("A storefront with a large wooden sign that reads", + "\"DIFFUSER\" in bold carved letters"), seed = 42) + # Or through the common dispatcher txt2img("a lighthouse at dusk", model_name = "flux2") ``` @@ -171,6 +177,7 @@ Currently supported models: - Stable Diffusion XL (SDXL) - FLUX.1-schnell (12B, 4-step distilled) - FLUX.2 Klein 4B (4-step distilled) +- Z-Image-Turbo (6B, 8-step distilled, text rendering) - LTX-2.3 Video (with audio) ### Choosing an image model: SDXL vs FLUX.2 @@ -182,19 +189,19 @@ Same prompt, same seed, 1024x1024, measured on an RTX 5060 Ti 16GB: | Model | Settings | Load | Warm generation | Peak VRAM | |---|---|---|---|---| -| SDXL | 50 steps, CFG 7.5 | 45 s | **20 s** | 12.7 GB | -| FLUX.2 Klein 4B | 4 steps, guidance-free | 32 s | 48 s | **8.2 GB** | +| SDXL | 50 steps, CFG 7.5 | 45 s | 20 s | **12.7 GB** | +| FLUX.2 Klein 4B | 4 steps, guidance-free | 32 s | **13 s** | 12.5 GB | -SDXL is over twice as fast per image, but FLUX.2's prompt adherence -and coherence are in a different class — SDXL melts the speaker -stacks and turns the wall records into neon smears, while FLUX.2 -draws the studio you asked for (SDXL left, FLUX.2 right): +FLUX.2 is now faster per image (since the allocator gc-gate fix), and +its prompt adherence and coherence are in a different class — SDXL +melts the speaker stacks and turns the wall records into neon smears, +while FLUX.2 draws the studio you asked for (SDXL left, FLUX.2 right): ![SDXL vs FLUX.2, same prompt and seed](man/figures/sdxl_vs_flux2_radio_studio.jpg) -Rule of thumb: reach for FLUX.2 unless you need images in bulk and -fast more than you need them right. Note FLUX.2 Klein is -guidance-free, so negative prompts do not apply. +Rule of thumb: reach for FLUX.2. Note FLUX.2 Klein is guidance-free, +so negative prompts do not apply. When the image needs legible text +(signs, posters, labels), reach for Z-Image-Turbo instead. ### Downloading Models diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md index 6c2c651..fb53187 100644 --- a/inst/REFERENCES.md +++ b/inst/REFERENCES.md @@ -40,6 +40,16 @@ module; this file documents the actual lineage, idea by idea. | Qwen2 byte-level BPE tokenization (GPT-2 byte table, GPT-4-style split regex, rank-based merges) | GPT-2 BPE (Radford et al. 2019; openai/gpt-2 `encoder.py`, MIT); format facts from HuggingFace tokenizers documentation | | Weights | black-forest-labs/FLUX.2-klein-4B (Apache-2.0, ungated; downloaded by the user, never redistributed) | +## Z-Image Turbo + +| What | Source | +|---|---| +| Single-stream DiT (sandwich RMSNorms, scale/gate-only tanh modulation, noise/context refiner stacks, SwiGLU with separate gate, 3-axis RoPE theta 256, SEQ_MULTI_OF padding with learned pad tokens) | Ported from HuggingFace **diffusers** (Apache-2.0): `models/transformers/transformer_z_image.py` | +| Pipeline flow: Qwen3 prompt encoding contract (thinking-enabled chat template, penultimate hidden state, mask-sliced captions), linspace sigma schedule, reversed-timestep + negated-output convention, static shift | diffusers `pipelines/z_image/pipeline_z_image.py` (Apache-2.0) | +| VAE decode | The FLUX.1 16-channel AutoencoderKL port, unchanged (the checkpoint's VAE config is the flux-dev autoencoder) | +| Qwen3-4B encoder + Qwen2 byte-level BPE tokenizer | Shared with the FLUX.2 klein port (see above); the tokenizer files are byte-identical between the two checkpoints | +| Weights | Tongyi-MAI/Z-Image-Turbo (Apache-2.0, ungated; downloaded by the user, never redistributed) | + ## Quantization | What | Source | diff --git a/inst/tinytest/fixtures/dit_zimage.safetensors b/inst/tinytest/fixtures/dit_zimage.safetensors new file mode 100644 index 0000000..410e72e Binary files /dev/null and b/inst/tinytest/fixtures/dit_zimage.safetensors differ diff --git a/inst/tinytest/fixtures/rope_zimage.safetensors b/inst/tinytest/fixtures/rope_zimage.safetensors new file mode 100644 index 0000000..241b40b Binary files /dev/null and b/inst/tinytest/fixtures/rope_zimage.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_model.safetensors b/inst/tinytest/fixtures/zimage_model.safetensors new file mode 100644 index 0000000..8b3fa6f Binary files /dev/null and b/inst/tinytest/fixtures/zimage_model.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_qwen.safetensors b/inst/tinytest/fixtures/zimage_qwen.safetensors new file mode 100644 index 0000000..9d76664 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_qwen.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_template_cases.json b/inst/tinytest/fixtures/zimage_template_cases.json new file mode 100644 index 0000000..9bad70c --- /dev/null +++ b/inst/tinytest/fixtures/zimage_template_cases.json @@ -0,0 +1,694 @@ +{ + "templated": [ + { + "text": "a photo of a cat", + "rendered": "<|im_start|>user\na photo of a cat<|im_end|>\n<|im_start|>assistant\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 64, + 6548, + 315, + 264, + 8251, + 151645, + 198, + 151644, + 77091, + 198, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "emoji 🦊 and 中文字符 mixed in", + "rendered": "<|im_start|>user\nemoji 🦊 and 中文字符 mixed in<|im_end|>\n<|im_start|>assistant\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 37523, + 11162, + 99, + 232, + 323, + 72858, + 16744, + 48391, + 9519, + 304, + 151645, + 198, + 151644, + 77091, + 198, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "An astronaut riding a horse on Mars, photorealistic", + "rendered": "<|im_start|>user\nAn astronaut riding a horse on Mars, photorealistic<|im_end|>\n<|im_start|>assistant\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 2082, + 46633, + 19837, + 264, + 15223, + 389, + 21048, + 11, + 4503, + 89768, + 4532, + 151645, + 198, + 151644, + 77091, + 198, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。", + "rendered": "<|im_start|>user\n一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。<|im_end|>\n<|im_start|>assistant\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 107814, + 17714, + 101599, + 2073, + 66078, + 48921, + 12881, + 57, + 12, + 29926, + 9285, + 1511, + 4677, + 10429, + 854, + 9370, + 73345, + 70500, + 9370, + 102343, + 108755, + 1773, + 151645, + 198, + 151644, + 77091, + 198, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + }, + { + "text": "", + "rendered": "<|im_start|>user\n<|im_end|>\n<|im_start|>assistant\n", + "max_length": 64, + "ids": [ + 151644, + 872, + 198, + 151645, + 198, + 151644, + 77091, + 198, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643, + 151643 + ], + "mask": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0 + ] + } + ], + "meta": { + "pad_token": "<|endoftext|>", + "pad_token_id": 151643, + "padding_side": "right" + } +} \ No newline at end of file diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/config.json b/inst/tinytest/fixtures/zimage_tiny_ckpt/config.json new file mode 100644 index 0000000..c2ffb58 --- /dev/null +++ b/inst/tinytest/fixtures/zimage_tiny_ckpt/config.json @@ -0,0 +1,32 @@ +{ + "_class_name": "ZImageTransformer2DModel", + "_diffusers_version": "0.39.0.dev0", + "all_f_patch_size": [ + 1 + ], + "all_patch_size": [ + 2 + ], + "axes_dims": [ + 8, + 8, + 8 + ], + "axes_lens": [ + 128, + 32, + 32 + ], + "cap_feat_dim": 24, + "dim": 48, + "in_channels": 4, + "n_heads": 2, + "n_kv_heads": 2, + "n_layers": 2, + "n_refiner_layers": 1, + "norm_eps": 1e-05, + "qk_norm": true, + "rope_theta": 256.0, + "siglip_feat_dim": null, + "t_scale": 1000.0 +} diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00001-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00001-of-00008.safetensors new file mode 100644 index 0000000..0d0cc35 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00001-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00002-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00002-of-00008.safetensors new file mode 100644 index 0000000..98af620 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00002-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00003-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00003-of-00008.safetensors new file mode 100644 index 0000000..416a130 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00003-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00004-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00004-of-00008.safetensors new file mode 100644 index 0000000..12fd03a Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00004-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00005-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00005-of-00008.safetensors new file mode 100644 index 0000000..8b1374e Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00005-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00006-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00006-of-00008.safetensors new file mode 100644 index 0000000..f6b7361 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00006-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00007-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00007-of-00008.safetensors new file mode 100644 index 0000000..06db1e7 Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00007-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00008-of-00008.safetensors b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00008-of-00008.safetensors new file mode 100644 index 0000000..7e2063a Binary files /dev/null and b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model-00008-of-00008.safetensors differ diff --git a/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json new file mode 100644 index 0000000..8cf332c --- /dev/null +++ b/inst/tinytest/fixtures/zimage_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json @@ -0,0 +1,80 @@ +{ + "metadata": { + "total_size": 1829664 + }, + "weight_map": { + "all_final_layer.2-1.adaLN_modulation.1.bias": "diffusion_pytorch_model-00001-of-00008.safetensors", + "all_final_layer.2-1.adaLN_modulation.1.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "all_final_layer.2-1.linear.bias": "diffusion_pytorch_model-00001-of-00008.safetensors", + "all_final_layer.2-1.linear.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "all_x_embedder.2-1.bias": "diffusion_pytorch_model-00001-of-00008.safetensors", + "all_x_embedder.2-1.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "cap_embedder.0.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "cap_embedder.1.bias": "diffusion_pytorch_model-00005-of-00008.safetensors", + "cap_embedder.1.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "cap_pad_token": "diffusion_pytorch_model-00001-of-00008.safetensors", + "context_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "context_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "context_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.0.attention.norm_k.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.attention.norm_q.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.attention.to_k.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.attention.to_out.0.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.attention.to_q.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.attention.to_v.weight": "diffusion_pytorch_model-00005-of-00008.safetensors", + "layers.0.attention_norm1.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.attention_norm2.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.feed_forward.w1.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.feed_forward.w2.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.feed_forward.w3.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.ffn_norm1.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.0.ffn_norm2.weight": "diffusion_pytorch_model-00006-of-00008.safetensors", + "layers.1.adaLN_modulation.0.bias": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.adaLN_modulation.0.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.attention.norm_k.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention.norm_q.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention.to_k.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention.to_out.0.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention.to_q.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention.to_v.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.attention_norm1.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.attention_norm2.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.feed_forward.w1.weight": "diffusion_pytorch_model-00007-of-00008.safetensors", + "layers.1.feed_forward.w2.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.feed_forward.w3.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.ffn_norm1.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "layers.1.ffn_norm2.weight": "diffusion_pytorch_model-00008-of-00008.safetensors", + "noise_refiner.0.adaLN_modulation.0.bias": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.adaLN_modulation.0.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.attention.norm_k.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention.norm_q.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention.to_k.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention.to_out.0.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention.to_q.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention.to_v.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.attention_norm1.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.attention_norm2.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.feed_forward.w1.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.feed_forward.w2.weight": "diffusion_pytorch_model-00001-of-00008.safetensors", + "noise_refiner.0.feed_forward.w3.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.ffn_norm1.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "noise_refiner.0.ffn_norm2.weight": "diffusion_pytorch_model-00002-of-00008.safetensors", + "t_embedder.mlp.0.bias": "diffusion_pytorch_model-00005-of-00008.safetensors", + "t_embedder.mlp.0.weight": "diffusion_pytorch_model-00003-of-00008.safetensors", + "t_embedder.mlp.2.bias": "diffusion_pytorch_model-00005-of-00008.safetensors", + "t_embedder.mlp.2.weight": "diffusion_pytorch_model-00004-of-00008.safetensors", + "x_pad_token": "diffusion_pytorch_model-00001-of-00008.safetensors" + } +} diff --git a/inst/tinytest/test_dit_zimage.R b/inst/tinytest/test_dit_zimage.R new file mode 100644 index 0000000..5a8c7df --- /dev/null +++ b/inst/tinytest/test_dit_zimage.R @@ -0,0 +1,103 @@ +# Parity tests for the Z-Image transformer blocks against diffusers +# reference fixtures (generated by tools/gen_fixtures_zimage_dit.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "dit_zimage.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/dit_zimage.safetensors" +if (!file.exists(fixture_path)) exit_file("zimage dit fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +load_named_weights <- function(module, weights) { + dests <- c(module$named_parameters(), module$named_buffers()) + missing_dest <- setdiff(names(weights), names(dests)) + if (length(missing_dest)) { + stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", ")) + } + unfilled <- setdiff(names(dests), names(weights)) + if (length(unfilled)) { + stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", ")) + } + torch::with_no_grad({ + for (name in names(weights)) dests[[name]]$copy_(weights[[name]]) + }) + invisible(module) +} + +subset_prefix <- function(fx, prefix) { + keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE) + out <- fx[keys] + names(out) <- sub(paste0("^", prefix, "\\."), "", keys) + out +} + +freqs <- list(fx$freqs_cos, fx$freqs_sin) + +# --- modulated block (noise refiner / main trunk) -------------------------------- + +blk <- zimage_block(dim = 32L, n_heads = 2L, modulation = TRUE) +load_named_weights(blk, subset_prefix(fx, "mod")) +blk$eval() +out <- torch::with_no_grad(blk(fx$x, freqs, adaln_input = fx$adaln)) +expect_equal(out$shape, c(1L, 24L, 32L)) +expect_true(max_abs_diff(out, fx$mod_out) < 1e-5) + +# --- unmodulated block (context refiner) ------------------------------------------ + +ublk <- zimage_block(dim = 32L, n_heads = 2L, modulation = FALSE) +load_named_weights(ublk, subset_prefix(fx, "unmod")) +ublk$eval() +uout <- torch::with_no_grad(ublk(fx$x, freqs)) +expect_true(max_abs_diff(uout, fx$unmod_out) < 1e-5) + +# The unmodulated block must carry no adaLN weights +expect_false(any(grepl("adaLN", names(ublk$named_parameters())))) + +# --- chunked attention matches unchunked ------------------------------------------- + +out_chunked <- torch::with_no_grad( + blk(fx$x, freqs, adaln_input = fx$adaln, chunk_size = 7L)) +expect_true(max_abs_diff(out_chunked, out) < 1e-6) + +# --- final layer ------------------------------------------------------------------- + +fin <- zimage_final_layer(hidden_size = 32L, out_channels = 64L) +load_named_weights(fin, subset_prefix(fx, "final")) +fin$eval() +fout <- torch::with_no_grad(fin(fx$x, fx$adaln)) +expect_equal(fout$shape, c(1L, 24L, 64L)) +expect_true(max_abs_diff(fout, fx$final_out) < 1e-5) + +# --- timestep embedder -------------------------------------------------------------- + +temb <- zimage_t_embedder(out_size = 32L, mid_size = 48L) +load_named_weights(temb, subset_prefix(fx, "temb")) +temb$eval() +tout <- torch::with_no_grad(temb(fx$temb_in)) +expect_true(max_abs_diff(tout, fx$temb_out) < 1e-5) + +# --- cap embedder (Sequential RMSNorm + Linear, keys 0/1) ---------------------------- + +cap <- torch::nn_sequential( + diffuseR:::ltx23_rms_norm(24L, eps = 1e-5), + torch::nn_linear(24L, 32L, bias = TRUE) +) +load_named_weights(cap, subset_prefix(fx, "cap")) +cap$eval() +cout <- torch::with_no_grad(cap(fx$cap_in)) +expect_true(max_abs_diff(cout, fx$cap_out) < 1e-6) diff --git a/inst/tinytest/test_quantize_zimage.R b/inst/tinytest/test_quantize_zimage.R new file mode 100644 index 0000000..75b96de --- /dev/null +++ b/inst/tinytest/test_quantize_zimage.R @@ -0,0 +1,114 @@ +# Z-Image quantization round trip on the tiny sharded checkpoint: +# family auto-detection via config _class_name, cast census, NF4 and +# fp8 (streamed + resident) loads. Everything runs on CPU. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +ckpt_dir <- system.file("tinytest", "fixtures", "zimage_tiny_ckpt", + package = "diffuseR") +if (ckpt_dir == "") ckpt_dir <- "fixtures/zimage_tiny_ckpt" +if (!dir.exists(ckpt_dir)) exit_file("zimage tiny checkpoint missing") + +fixture_path <- system.file("tinytest", "fixtures", "zimage_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/zimage_model.safetensors" +if (!file.exists(fixture_path)) exit_file("zimage model fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} +cosine_sim <- function(a, b) { + a <- a$to(dtype = torch::torch_float32())$flatten() + b <- b$to(dtype = torch::torch_float32())$flatten() + as.numeric(torch::torch_dot(a, b) / (a$norm() * b$norm())) +} + +# --- family detection + cast census ------------------------------------------------ +# Tiny config: (2 layers + 1 noise_refiner + 1 context_refiner) x 7 = 28 +# cast weights. (Full Turbo: 34 blocks x 7 = 238.) + +ckpt <- flux_open_checkpoint(ckpt_dir) +expect_equal(diffuseR:::.flux_family(ckpt$config), "zimage") +expect_equal(sum(zimage_is_quant_key(ckpt$keys)), 28L) + +# The adaLN modulation linears and embedders stay resident +expect_false(any(zimage_is_quant_key( + c("layers.0.adaLN_modulation.0.weight", "cap_embedder.1.weight", + "all_x_embedder.2-1.weight", "t_embedder.mlp.0.weight", + "x_pad_token", "layers.0.attention_norm1.weight") +))) + +# --- full-precision load through family dispatch ------------------------------------ + +model <- flux_load_transformer(ckpt, device = "cpu", dtype = "float32", + verbose = FALSE) +out_full <- torch::with_no_grad(model(fx$x, fx$t, fx$cap)) +expect_true(max_abs_diff(out_full, fx$out) < 1e-4) + +# --- NF4 round trip ------------------------------------------------------------------ + +nf4_dir <- file.path(tempdir(), "zimage-tiny-nf4") +unlink(nf4_dir, recursive = TRUE) +manifest <- flux_quantize(ckpt_dir, output_dir = nf4_dir, format = "nf4", + verbose = FALSE) +expect_equal(manifest$cast, 28L) +expect_true(grepl("^zimage-turbo-nf4", manifest$shards[[1]])) + +model_nf4 <- flux_load_transformer(flux_open_quantized(nf4_dir), + device = "cpu", verbose = FALSE) +out_nf4 <- torch::with_no_grad(model_nf4( + fx$x$to(dtype = torch::torch_bfloat16()), fx$t, + fx$cap$to(dtype = torch::torch_bfloat16()) +)) +expect_true(cosine_sim(out_nf4, out_full) > 0.98) +ltx23_release_dequant_buffers() + +# --- fp8 round trips (streamed and resident) ------------------------------------------ + +f8_ok <- tryCatch({ + x <- torch::torch_randn(2, 2)$to(dtype = torch::torch_float8_e4m3fn()) + tmp <- tempfile(fileext = ".safetensors") + safetensors::safe_save_file(list(w = x), tmp) + y <- safetensors::safe_load_file(tmp, framework = "torch") + unlink(tmp) + TRUE +}, error = function(e) FALSE) + +if (f8_ok) { + fp8_dir <- file.path(tempdir(), "zimage-tiny-fp8") + unlink(fp8_dir, recursive = TRUE) + manifest8 <- flux_quantize(ckpt_dir, output_dir = fp8_dir, format = "fp8", + verbose = FALSE) + expect_equal(manifest8$cast, 28L) + + model_fp8 <- flux_load_transformer(flux_open_quantized(fp8_dir), + device = "cpu", pin = FALSE, verbose = FALSE) + out_fp8 <- torch::with_no_grad(model_fp8( + fx$x$to(dtype = torch::torch_bfloat16()), fx$t, + fx$cap$to(dtype = torch::torch_bfloat16()) + )) + expect_true(cosine_sim(out_fp8, out_full) > 0.99) + + # Resident variant: weights moved to the compute device (CPU here just + # exercises the walker), forward unchanged + model_res <- flux_load_transformer(flux_open_quantized(fp8_dir), + device = "cpu", pin = FALSE, fp8_resident = TRUE, verbose = FALSE) + out_res <- torch::with_no_grad(model_res( + fx$x$to(dtype = torch::torch_bfloat16()), fx$t, + fx$cap$to(dtype = torch::torch_bfloat16()) + )) + expect_true(max_abs_diff(out_res, out_fp8) == 0) +} + +options(diffuseR.block_gc = NULL) diff --git a/inst/tinytest/test_rope_zimage.R b/inst/tinytest/test_rope_zimage.R new file mode 100644 index 0000000..4a9b881 --- /dev/null +++ b/inst/tinytest/test_rope_zimage.R @@ -0,0 +1,100 @@ +# Parity tests for the Z-Image position ids, 3-axis RoPE, patchify chain, +# timestep sinusoid, and static-shift scheduler against diffusers +# reference fixtures (generated by tools/gen_fixtures_zimage.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "rope_zimage.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/rope_zimage.safetensors" +if (!file.exists(fixture_path)) exit_file("zimage rope fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# --- padding lengths ------------------------------------------------------------ + +pad_len <- diffuseR:::zimage_pad_len +expect_equal(pad_len(37L), 27L) +expect_equal(pad_len(60L), 4L) +expect_equal(pad_len(64L), 0L) +expect_equal(pad_len(4096L), 0L) + +# --- position ids --------------------------------------------------------------- + +# Caption: 37 tokens pad to 64; ramp 1..64 on axis 1 built over the +# padded length +cap_ids <- zimage_cap_pos_ids(37L + pad_len(37L)) +expect_equal(cap_ids$shape, c(64L, 3L)) +expect_true(max_abs_diff(cap_ids, fx$cap_pos_ids) == 0) + +# Image: 6x10 grid = 60 tokens pad to 64; axis-1 start just past the +# padded caption, pads at the origin +img_ids <- zimage_img_pos_ids(6L, 10L, start0 = 65L) +expect_equal(img_ids$shape, c(64L, 3L)) +expect_true(max_abs_diff(img_ids, fx$img_pos_ids) == 0) + +# --- RoPE frequencies (theta 256, f32-angle cast) -------------------------------- + +freqs <- zimage_pos_embed(fx$rope_ids) +expect_equal(freqs[[1]]$shape, c(9L, 128L)) +expect_true(max_abs_diff(freqs[[1]], fx$rope_cos) < 1e-6) +expect_true(max_abs_diff(freqs[[2]], fx$rope_sin) < 1e-6) + +# Position (0,0,0) is the identity rotation +expect_true(max_abs_diff(freqs[[1]][9, ], + torch::torch_ones(128L)) == 0) +expect_true(max_abs_diff(freqs[[2]][9, ], + torch::torch_zeros(128L)) == 0) + +# --- RoPE application (reuses the FLUX interleaved apply) ------------------------- + +# Fixture x is [B, S, H, D] (processor layout); flux_apply_rotary_emb +# takes [B, H, S, D] +x_bhsd <- fx$rope_x$permute(c(1L, 3L, 2L, 4L)) +out <- flux_apply_rotary_emb(x_bhsd, freqs) +expect_true(max_abs_diff(out, fx$rope_out$permute(c(1L, 3L, 2L, 4L))) < 1e-5) + +# --- patchify / unpatchify -------------------------------------------------------- + +tokens <- zimage_patchify(fx$img) +expect_equal(tokens$shape, c(60L, 64L)) +expect_true(max_abs_diff(tokens, fx$img_tokens_padded[1:60, ]) == 0) + +unpat <- zimage_unpatchify(fx$unpat_tokens, c(1L, 12L, 20L)) +expect_equal(unpat$shape, c(16L, 1L, 12L, 20L)) +expect_true(max_abs_diff(unpat, fx$unpat_out) == 0) + +# Roundtrip identity +rt <- zimage_unpatchify(zimage_patchify(fx$img), c(1L, 12L, 20L)) +expect_true(max_abs_diff(rt, fx$img) == 0) + +# --- timestep sinusoid ------------------------------------------------------------- + +t_emb <- ltx23_get_timestep_embedding(fx$t_emb_in, 256L, + flip_sin_to_cos = TRUE, downscale_freq_shift = 0) +expect_true(max_abs_diff(t_emb, fx$t_emb_out) < 1e-5) + +# --- scheduler: static shift 3.0 on linspace(1, 1/N, N) --------------------------- + +for (n in c(4L, 8L)) { + sched <- flowmatch_scheduler_create(shift = 3.0, use_dynamic_shifting = FALSE) + sched <- flowmatch_set_timesteps(sched, num_inference_steps = n, + sigmas = seq(1, 1 / n, length.out = n)) + expect_true(max_abs_diff(sched$sigmas, + fx[[paste0("sched_sigmas_", n)]]) < 1e-6) + expect_true(max_abs_diff(sched$timesteps, + fx[[paste0("sched_timesteps_", n)]]) < 1e-3) +} diff --git a/inst/tinytest/test_zimage_qwen.R b/inst/tinytest/test_zimage_qwen.R new file mode 100644 index 0000000..a8ed0ea --- /dev/null +++ b/inst/tinytest/test_zimage_qwen.R @@ -0,0 +1,116 @@ +# Z-Image Qwen3 delta: the enable_thinking=TRUE chat template against +# shipped-tokenizer renders (tools/gen_zimage_qwen_template_cases.py) and +# the hidden_states[-2] + mask-slice convention against a tiny reference +# model (tools/gen_fixtures_zimage_qwen.py). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +# --- hidden_states[-2] + mask slicing (tiny reference model) ---------------------- + +fixture_path <- system.file("tinytest", "fixtures", "zimage_qwen.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/zimage_qwen.safetensors" +if (!file.exists(fixture_path)) exit_file("zimage qwen fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +sd <- fx[grep("^sd\\.", names(fx))] +names(sd) <- paste0("model.", sub("^sd\\.", "", names(sd))) + +enc <- qwen3_encoder( + vocab_size = 128L, hidden_size = 32L, intermediate_size = 64L, + num_hidden_layers = 4L, num_attention_heads = 4L, + num_key_value_heads = 2L, head_dim = 8L, rope_theta = 1e6, + rms_norm_eps = 1e-6 +) +dests <- c(enc$named_parameters(), enc$named_buffers()) +# The final norm never runs for intermediate hidden states; every +# checkpoint key must still land +expect_true(all(names(sd) %in% names(dests))) +torch::with_no_grad({ + for (name in names(sd)) dests[[name]]$copy_(sd[[name]]) +}) +enc$eval() + +ids <- torch::torch_tensor( + matrix(as.integer(as.array(fx$input_ids)) + 1L, nrow = 1), + dtype = torch::torch_long() +) +mask <- fx$attention_mask$to(dtype = torch::torch_long()) + +# hidden_states[-2] of a 4-layer model = state after layer 3 +states <- torch::with_no_grad(enc(ids, attention_mask = mask, out_layers = 3L)) +expect_true(max_abs_diff(states[[1]], fx$penult) < 1e-5) + +# Mask slice to the variable-length caption (right padding -> first n) +n_real <- as.integer(sum(as.array(fx$attention_mask))) +expect_equal(n_real, 9L) +sliced <- states[[1]][1, 1:n_real, ] +expect_true(max_abs_diff(sliced, fx$penult_sliced) < 1e-5) + +# --- chat template with enable_thinking = TRUE ------------------------------------- + +find_qwen_tokenizer <- function() { + p <- Sys.getenv("DIFFUSER_QWEN_TOKENIZER", "") + if (nzchar(p) && file.exists(p)) { + return(p) + } + if (requireNamespace("hfhub", quietly = TRUE)) { + for (repo in c( + "Tongyi-MAI/Z-Image-Turbo", + "black-forest-labs/FLUX.2-klein-4B" + )) { + p <- tryCatch( + suppressMessages(hfhub::hub_download( + repo, "tokenizer/tokenizer.json", local_files_only = TRUE + )), + error = function(e) "" + ) + if (nzchar(p) && file.exists(p)) { + return(p) + } + } + } + p <- "../../tools/cache/tokenizer_qwen.json" + if (file.exists(p)) { + return(p) + } + "" +} + +tok_path <- find_qwen_tokenizer() +if (!nzchar(tok_path)) exit_file("no Qwen tokenizer.json available") + +cases_path <- system.file("tinytest", "fixtures", "zimage_template_cases.json", + package = "diffuseR") +if (cases_path == "") cases_path <- "fixtures/zimage_template_cases.json" +if (!file.exists(cases_path)) exit_file("zimage template cases missing") + +cases <- jsonlite::fromJSON(cases_path, simplifyVector = FALSE) + +expect_equal(cases$meta$padding_side, "right") +expect_equal(cases$meta$pad_token_id, 151643L) + +tok <- qwen_bpe_tokenizer(tok_path) + +for (case in cases$templated) { + got <- encode_qwen(tok, case$text, max_length = case$max_length, + chat_template = TRUE, enable_thinking = TRUE) + expect_equal(got$input_ids[1, ], as.integer(unlist(case$ids)), + info = sprintf("ids for: %s", substr(case$text, 1, 40))) + expect_equal(got$attention_mask[1, ], as.integer(unlist(case$mask)), + info = sprintf("mask for: %s", substr(case$text, 1, 40))) +} diff --git a/inst/tinytest/test_zimage_transformer.R b/inst/tinytest/test_zimage_transformer.R new file mode 100644 index 0000000..9c193c9 --- /dev/null +++ b/inst/tinytest/test_zimage_transformer.R @@ -0,0 +1,54 @@ +# Full-forward parity test for the Z-Image transformer against a tiny +# random-init diffusers reference model (generated by +# tools/gen_fixtures_zimage_model.py). Includes a strict bidirectional +# key census. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +library(diffuseR) + +fixture_path <- system.file("tinytest", "fixtures", "zimage_model.safetensors", + package = "diffuseR") +if (fixture_path == "") fixture_path <- "fixtures/zimage_model.safetensors" +if (!file.exists(fixture_path)) exit_file("zimage model fixtures missing") + +fx <- safetensors::safe_load_file(fixture_path, framework = "torch") + +max_abs_diff <- function(a, b) { + as.numeric(torch::torch_max(torch::torch_abs( + a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32()) + ))) +} + +sd <- fx[grep("^sd\\.", names(fx))] +names(sd) <- sub("^sd\\.", "", names(sd)) + +model <- zimage_transformer( + in_channels = 4L, dim = 48L, n_layers = 2L, n_refiner_layers = 1L, + n_heads = 2L, cap_feat_dim = 24L, axes_dims = c(8L, 8L, 8L) +) + +# Strict bidirectional key census: every checkpoint key lands, every +# module parameter fills, and the context refiner carries no adaLN +dests <- c(model$named_parameters(), model$named_buffers()) +expect_equal(sort(names(dests)), sort(names(sd))) +expect_false(any(grepl("context_refiner.*adaLN", names(dests)))) +expect_true(any(grepl("noise_refiner.*adaLN", names(dests)))) + +torch::with_no_grad({ + for (name in names(sd)) dests[[name]]$copy_(sd[[name]]) +}) +model$eval() + +out <- torch::with_no_grad(model(fx$x, fx$t, fx$cap)) +expect_equal(out$shape, c(4L, 1L, 12L, 20L)) +expect_true(max_abs_diff(out, fx$out) < 1e-4) + +# Chunked attention path matches +out_chunked <- torch::with_no_grad(model(fx$x, fx$t, fx$cap, chunk_size = 16L)) +expect_true(max_abs_diff(out_chunked, out) < 1e-5) diff --git a/man/dit_zimage_modules.Rd b/man/dit_zimage_modules.Rd new file mode 100644 index 0000000..f342f66 --- /dev/null +++ b/man/dit_zimage_modules.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{dit_zimage_modules} +\alias{dit_zimage_modules} +\title{Z-Image Transformer Block Modules} +\description{ +Fresh R port of the Z-Image DiT building blocks from the diffusers +reference (Apache-2.0, +src/diffusers/models/transformers/transformer_z_image.py). Z-Image is +a single-stream DiT: text and image tokens share one sequence and one +set of block weights. Each block uses sandwich RMSNorms (a learned +norm before AND after both the attention and the feed-forward) and a +scale/gate-only modulation — four chunks (scale_msa, gate_msa, +scale_mlp, gate_mlp), no shift, gates tanh-squashed, scales 1 + x. +The attention is plain joint self-attention, so the FLUX attention +module is reused with bias = FALSE and eps = 1e-5. +} diff --git a/man/download_zimage.Rd b/man/download_zimage.Rd new file mode 100644 index 0000000..23b958f --- /dev/null +++ b/man/download_zimage.Rd @@ -0,0 +1,11 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_zimage} +\alias{download_zimage} +\title{Download and Prepare Z-Image-Turbo Weights} +\description{ +Downloads Z-Image-Turbo from HuggingFace (Apache-2.0, ungated) and +quantizes the 6B transformer to a local fp8 (~6.3 GB) or NF4 +(~3.6 GB) artifact. The checkpoint ships the transformer in float32 +(24.6 GB), so the one-time quantize saves a lot of disk and load +time. +} diff --git a/man/download_zimage_turbo.Rd b/man/download_zimage_turbo.Rd new file mode 100644 index 0000000..2cb8cd4 --- /dev/null +++ b/man/download_zimage_turbo.Rd @@ -0,0 +1,32 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_zimage_turbo} +\alias{download_zimage_turbo} +\title{Download Z-Image-Turbo and build the quantized artifact} +\usage{ +download_zimage_turbo(quantize = TRUE, precision = c("fp8", "nf4"), + output_dir = NULL, text_encoders = TRUE, verbose = TRUE) +} +\arguments{ +\item{quantize}{Logical. Build the quantized artifact.} + +\item{precision}{"fp8" (~6.3 GB, GPU-resident; near-bf16 quality) or +"nf4" (~3.6 GB).} + +\item{output_dir}{Directory for the quantized artifact.} + +\item{text_encoders}{Logical. Also fetch the Qwen3-4B text encoder, +tokenizer, VAE, and scheduler config (~8.2 GB).} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, a list with \code{transformer_dir}, + \code{artifact_dir}, and \code{support} (named file paths). +} +\description{ +Skips work already done: a valid quantized manifest short-circuits +the transformer download; cached files are not re-fetched. No token +is needed (the repo is ungated). The float32 transformer source +(~24.6 GB in the HuggingFace cache) may be deleted after +quantization. +} diff --git a/man/encode_qwen.Rd b/man/encode_qwen.Rd index c4d784e..ae78b95 100644 --- a/man/encode_qwen.Rd +++ b/man/encode_qwen.Rd @@ -3,7 +3,8 @@ \alias{encode_qwen} \title{Encode prompts with the Qwen tokenizer} \usage{ -encode_qwen(tokenizer, texts, max_length = 512L, chat_template = TRUE) +encode_qwen(tokenizer, texts, max_length = 512L, chat_template = TRUE, + enable_thinking = FALSE) } \arguments{ \item{tokenizer}{A \code{\link{qwen_bpe_tokenizer}}.} @@ -14,6 +15,9 @@ encode_qwen(tokenizer, texts, max_length = 512L, chat_template = TRUE) for no truncation/padding.} \item{chat_template}{Logical. Wrap in the Qwen3 chat template.} + +\item{enable_thinking}{Logical. Leave the model's thinking enabled +(no empty think block). Default FALSE.} } \value{ List with \code{input_ids} and \code{attention_mask} integer @@ -21,9 +25,11 @@ List with \code{input_ids} and \code{attention_mask} integer \code{max_length} is NULL). Ids are 0-based. } \description{ -With \code{chat_template = TRUE} (the FLUX.2 klein pipeline behavior) -each prompt is wrapped as a single user turn with the generation -prompt and a disabled thinking block, matching -\code{apply_chat_template(..., add_generation_prompt = TRUE, -enable_thinking = FALSE)}. Right-pads with \code{<|endoftext|>}. +With \code{chat_template = TRUE} each prompt is wrapped as a single +user turn with the generation prompt, matching +\code{apply_chat_template(..., add_generation_prompt = TRUE)}. With +\code{enable_thinking = FALSE} (the FLUX.2 klein pipeline behavior) +the template closes with an empty thinking block; with +\code{enable_thinking = TRUE} (the Z-Image pipeline behavior) it ends +at the assistant turn. Right-pads with \code{<|endoftext|>}. } diff --git a/man/rope_zimage.Rd b/man/rope_zimage.Rd new file mode 100644 index 0000000..20a96cc --- /dev/null +++ b/man/rope_zimage.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{rope_zimage} +\alias{rope_zimage} +\title{Z-Image Rotary Positional Embeddings and Patchify Helpers} +\description{ +Fresh R port of the Z-Image position scheme from the diffusers +reference (Apache-2.0, +src/diffusers/models/transformers/transformer_z_image.py RopeEmbedder, +create_coordinate_grid, _patchify_image, _pad_with_ids, unpatchify). +Z-Image uses 3-axis interleaved RoPE with theta 256; frequencies are +built in float64 but the angles are cast to float32 before cos/sin +(torch.polar on a .float() tensor), which differs measurably from the +FLUX convention at large positions. Every sub-sequence is padded to a +multiple of 32 (SEQ_MULTI_OF); caption positions are a 1-based ramp on +axis 1 built over the padded length, image positions sit on axes 2/3 +with axis 1 offset just past the caption. +} diff --git a/man/txt2img.Rd b/man/txt2img.Rd index b0f5628..1557468 100644 --- a/man/txt2img.Rd +++ b/man/txt2img.Rd @@ -3,7 +3,7 @@ \alias{txt2img} \title{Generate an image from a text prompt using a diffusion pipeline} \usage{ -txt2img(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2"), ...) +txt2img(prompt, model_name = c("sd21", "sdxl", "flux1", "flux2", "zimage"), ...) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} diff --git a/man/txt2img_zimage.Rd b/man/txt2img_zimage.Rd new file mode 100644 index 0000000..be0d6ea --- /dev/null +++ b/man/txt2img_zimage.Rd @@ -0,0 +1,49 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{txt2img_zimage} +\alias{txt2img_zimage} +\title{Generate an image with Z-Image-Turbo} +\usage{ +txt2img_zimage(prompt, pipeline = NULL, width = 1024L, height = 1024L, + num_inference_steps = 8L, max_sequence_length = 512L, + seed = NULL, prompt_embeds = NULL, save_file = TRUE, + filename = NULL, verbose = TRUE, ...) +} +\arguments{ +\item{prompt}{Character. The prompt.} + +\item{pipeline}{A \code{zimage_pipeline} from +\code{\link{zimage_load_pipeline}}; NULL loads one (passing +\code{...} through).} + +\item{num_inference_steps}{Integer. Denoising steps (Turbo: 8).} + +\item{max_sequence_length}{Integer. Qwen3 token length (512).} + +\item{seed}{Integer or NULL. Latents are drawn on the CPU, so a seed +matches a Python diffusers run with a CPU generator.} + +\item{prompt_embeds}{Optional precomputed [L, 2560] caption +embeddings (valid tokens only).} + +\item{save_file}{Logical. Write a PNG.} + +\item{filename}{Output path (default derived from the prompt).} + +\item{verbose}{Logical.} + +\item{...}{Passed to \code{\link{zimage_load_pipeline}} when +\code{pipeline} is NULL.} + +\item{width,height}{Integers, divisible by 16.} +} +\value{ +Invisibly, \code{list(image, metadata)} where \code{image} is + an [H, W, 3] array in [0, 1]. +} +\description{ +Guidance-distilled text-to-image (8 steps, no CFG): Qwen3-4B prompt +encoding (thinking-enabled chat template, penultimate hidden state), +FlowMatch denoising with the reversed-timestep convention, and +16-channel VAE decode. Strong at legible text rendering, English and +Chinese both. +} diff --git a/man/zimage_block.Rd b/man/zimage_block.Rd new file mode 100644 index 0000000..04f67a1 --- /dev/null +++ b/man/zimage_block.Rd @@ -0,0 +1,23 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_block} +\alias{zimage_block} +\title{Z-Image transformer block} +\usage{ +zimage_block(dim, n_heads, norm_eps = 1e-05, modulation = TRUE) +} +\arguments{ +\item{dim}{Integer. Model width.} + +\item{n_heads}{Integer. Attention heads; head dim is dim / n_heads.} + +\item{norm_eps}{Numeric. RMSNorm epsilon. Default 1e-5.} + +\item{modulation}{Logical. Whether the block is timestep-modulated.} +} +\description{ +Sandwich-norm residual block shared by the noise refiner, the context +refiner and the main trunk. With \code{modulation = TRUE} the block +carries an adaLN linear producing (scale_msa, gate_msa, scale_mlp, +gate_mlp); the context refiner uses \code{modulation = FALSE} and has +no adaLN weights at all. +} diff --git a/man/zimage_cap_pos_ids.Rd b/man/zimage_cap_pos_ids.Rd new file mode 100644 index 0000000..abf15d2 --- /dev/null +++ b/man/zimage_cap_pos_ids.Rd @@ -0,0 +1,23 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_cap_pos_ids} +\alias{zimage_cap_pos_ids} +\title{Build Z-Image caption position ids} +\usage{ +zimage_cap_pos_ids(cap_padded_len, device = "cpu") +} +\arguments{ +\item{cap_padded_len}{Integer. Caption length after padding to a +multiple of 32.} + +\item{device}{Device for the resulting tensor.} +} +\value{ +Float tensor of shape [cap_padded_len, 3]. +} +\description{ +Caption tokens ramp 1..cap_padded_len on the first axis (axes 2 and 3 +zero). The reference builds the grid over the already-padded length, +so pad tokens continue the ramp rather than sitting at the origin +(the (0,0,0) pad ids it also emits are truncated away in +_prepare_sequence and never reach RoPE). +} diff --git a/man/zimage_feed_forward.Rd b/man/zimage_feed_forward.Rd new file mode 100644 index 0000000..7c49d8f --- /dev/null +++ b/man/zimage_feed_forward.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_feed_forward} +\alias{zimage_feed_forward} +\title{Z-Image feed-forward (SwiGLU with separate gate weights)} +\usage{ +zimage_feed_forward(dim, hidden_dim) +} +\arguments{ +\item{dim}{Integer. Model width.} + +\item{hidden_dim}{Integer. Hidden width.} +} +\description{ +w2(silu(w1(x)) * w3(x)) with all three linears bias-free. The hidden +width is int(dim / 3 * 8). +} diff --git a/man/zimage_final_layer.Rd b/man/zimage_final_layer.Rd new file mode 100644 index 0000000..9642f70 --- /dev/null +++ b/man/zimage_final_layer.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_final_layer} +\alias{zimage_final_layer} +\title{Z-Image final layer} +\usage{ +zimage_final_layer(hidden_size, out_channels) +} +\arguments{ +\item{hidden_size}{Integer. Model width.} + +\item{out_channels}{Integer. Patch output dim +(patch^2 * f_patch * latent channels).} +} +\description{ +Parameterless LayerNorm scaled by 1 + adaLN(c) (scale only, no +shift), then the token-to-patch projection. +} diff --git a/man/zimage_img_pos_ids.Rd b/man/zimage_img_pos_ids.Rd new file mode 100644 index 0000000..59ccec1 --- /dev/null +++ b/man/zimage_img_pos_ids.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_img_pos_ids} +\alias{zimage_img_pos_ids} +\title{Build Z-Image latent image position ids} +\usage{ +zimage_img_pos_ids(h_tokens, w_tokens, start0, f_tokens = 1L, device = "cpu") +} +\arguments{ +\item{h_tokens}{Integer. Token grid height (latent height / patch).} + +\item{w_tokens}{Integer. Token grid width (latent width / patch).} + +\item{start0}{Integer. First-axis start, cap_padded_len + 1.} + +\item{f_tokens}{Integer. Token grid frames; 1 for txt2img.} + +\item{device}{Device for the resulting tensor.} +} +\value{ +Float tensor of shape [padded token count, 3]. +} +\description{ +Image tokens use axis 1 for the frame index offset past the caption +(start0 = cap_padded_len + 1), axis 2 for the token row and axis 3 for +the token column. Trailing pad tokens (token count not a multiple of +32) sit at (0, 0, 0). Reference: patchify_and_embed / _pad_with_ids. +} diff --git a/man/zimage_is_quant_key.Rd b/man/zimage_is_quant_key.Rd new file mode 100644 index 0000000..0257ab0 --- /dev/null +++ b/man/zimage_is_quant_key.Rd @@ -0,0 +1,16 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_is_quant_key} +\alias{zimage_is_quant_key} +\title{Test whether a Z-Image key is in the quantization cast set} +\usage{ +zimage_is_quant_key(key) +} +\arguments{ +\item{key}{Character vector of parameter names (diffusers-style).} +} +\value{ +Logical vector. +} +\description{ +Test whether a Z-Image key is in the quantization cast set +} diff --git a/man/zimage_load_pipeline.Rd b/man/zimage_load_pipeline.Rd new file mode 100644 index 0000000..6438623 --- /dev/null +++ b/man/zimage_load_pipeline.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_load_pipeline} +\alias{zimage_load_pipeline} +\title{Load the Z-Image-Turbo pipeline} +\usage{ +zimage_load_pipeline(model_dir = NULL, device = "cuda", + precision = c("fp8", "nf4"), text_device = NULL, + attn_chunk = NULL, phase_offload = TRUE, verbose = TRUE) +} +\arguments{ +\item{model_dir}{Quantized artifact directory (default: the +\code{download_zimage_turbo} location for \code{precision}), or a +raw diffusers transformer directory.} + +\item{device}{Character. Compute device.} + +\item{precision}{"fp8" (default) or "nf4".} + +\item{text_device}{Device for the Qwen3 encoder (default: +\code{device}; it encodes in its own phase and offloads).} + +\item{attn_chunk}{Integer or NULL. Attention query-chunk override.} + +\item{phase_offload}{Logical. One GPU tenant per phase.} + +\item{verbose}{Logical.} +} +\value{ +A \code{zimage_pipeline} list. +} +\description{ +Loads the quantized transformer artifact plus the 16-channel VAE +decoder, Qwen3-4B text encoder, and tokenizer from the HuggingFace +cache populated by \code{\link{download_zimage_turbo}}. With fp8 +precision the ~6.3 GB transformer rides to the GPU per phase. +} diff --git a/man/zimage_pad_len.Rd b/man/zimage_pad_len.Rd new file mode 100644 index 0000000..fad01a3 --- /dev/null +++ b/man/zimage_pad_len.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_pad_len} +\alias{zimage_pad_len} +\title{Padding length to the next multiple of 32} +\usage{ +zimage_pad_len(n) +} +\arguments{ +\item{n}{Integer token count.} +} +\value{ +Integer pad length in [0, 31]. +} +\description{ +Padding length to the next multiple of 32 +} +\keyword{internal} diff --git a/man/zimage_patchify.Rd b/man/zimage_patchify.Rd new file mode 100644 index 0000000..f485e0f --- /dev/null +++ b/man/zimage_patchify.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_patchify} +\alias{zimage_patchify} +\title{Patchify a latent image to Z-Image tokens} +\usage{ +zimage_patchify(image, patch_size = 2L, f_patch_size = 1L) +} +\arguments{ +\item{image}{Tensor of shape [C, F, H, W].} + +\item{patch_size}{Integer spatial patch size. Default 2.} + +\item{f_patch_size}{Integer temporal patch size. Default 1.} +} +\value{ +Tensor of shape [num_tokens, patch_dim]. +} +\description{ +(C, F, H, W) -> [F/pF * H/p * W/p, pF * p * p * C], matching +_patchify_image. No padding is applied here. +} diff --git a/man/zimage_pos_embed.Rd b/man/zimage_pos_embed.Rd new file mode 100644 index 0000000..05364b0 --- /dev/null +++ b/man/zimage_pos_embed.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_pos_embed} +\alias{zimage_pos_embed} +\title{Compute Z-Image rotary frequencies from position ids} +\usage{ +zimage_pos_embed(ids, axes_dim = c(32L, 48L, 48L), theta = 256) +} +\arguments{ +\item{ids}{Tensor of shape [S, 3] from \code{zimage_cap_pos_ids} / +\code{zimage_img_pos_ids}.} + +\item{axes_dim}{Integer vector of per-axis rotary dims; must sum to +the attention head dim. Z-Image uses c(32, 48, 48).} + +\item{theta}{Numeric. RoPE base frequency. Z-Image uses 256.} +} +\value{ +List of two tensors (cos, sin), each [S, sum(axes_dim)], + float32, on the device of \code{ids}. +} +\description{ +Per-axis 1D rotary frequencies in the interleaved-real convention. +Frequencies and angles are built in float64, then the angles are cast +to float32 before cos/sin — matching the reference torch.polar call on +a .float() tensor. Output format matches \code{flux_pos_embed} so +\code{flux_apply_rotary_emb} applies unchanged. +} diff --git a/man/zimage_t_embedder.Rd b/man/zimage_t_embedder.Rd new file mode 100644 index 0000000..f8c3b73 --- /dev/null +++ b/man/zimage_t_embedder.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_t_embedder} +\alias{zimage_t_embedder} +\title{Z-Image timestep embedder} +\usage{ +zimage_t_embedder(out_size, mid_size = 1024L, freq_size = 256L) +} +\arguments{ +\item{out_size}{Integer. Output width, min(dim, 256).} + +\item{mid_size}{Integer. Hidden width. The full model uses 1024.} + +\item{freq_size}{Integer. Sinusoid width. Default 256.} +} +\description{ +256-dim cos-first sinusoid (computed in float32) through a +Linear-SiLU-Linear MLP. The model feeds t * t_scale with the +pipeline's t already in [0, 1]. +} diff --git a/man/zimage_transformer.Rd b/man/zimage_transformer.Rd new file mode 100644 index 0000000..af81f4b --- /dev/null +++ b/man/zimage_transformer.Rd @@ -0,0 +1,57 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_transformer} +\alias{zimage_transformer} +\title{Z-Image Transformer} +\usage{ +zimage_transformer(in_channels = 16L, dim = 3840L, n_layers = 30L, + n_refiner_layers = 2L, n_heads = 30L, norm_eps = 1e-05, + cap_feat_dim = 2560L, rope_theta = 256, t_scale = 1000, + axes_dims = c(32L, 48L, 48L), patch_size = 2L, + f_patch_size = 1L) +} +\arguments{ +\item{in_channels}{Integer. Latent channels. Default 16.} + +\item{dim}{Integer. Model width. Default 3840.} + +\item{n_layers}{Integer. Main trunk depth. Default 30.} + +\item{n_refiner_layers}{Integer. Refiner depth. Default 2.} + +\item{n_heads}{Integer. Attention heads. Default 30.} + +\item{norm_eps}{Numeric. RMSNorm epsilon. Default 1e-5.} + +\item{cap_feat_dim}{Integer. Caption embedding width. Default 2560.} + +\item{rope_theta}{Numeric. RoPE base frequency. Default 256.} + +\item{t_scale}{Numeric. Timestep scale. Default 1000.} + +\item{axes_dims}{Integer vector. Per-axis rotary dims. Default +c(32, 48, 48).} + +\item{patch_size}{Integer. Spatial patch size. Default 2.} + +\item{f_patch_size}{Integer. Temporal patch size. Default 1.} +} +\description{ +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). +} +\details{ +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). + +} diff --git a/man/zimage_unpatchify.Rd b/man/zimage_unpatchify.Rd new file mode 100644 index 0000000..9da8f85 --- /dev/null +++ b/man/zimage_unpatchify.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{zimage_unpatchify} +\alias{zimage_unpatchify} +\title{Unpatchify Z-Image tokens back to a latent image} +\usage{ +zimage_unpatchify(tokens, size, patch_size = 2L, f_patch_size = 1L, + out_channels = 16L) +} +\arguments{ +\item{tokens}{Tensor of shape [S, pF * p * p * C] with the image +tokens first.} + +\item{size}{Integer vector c(F, H, W) of the target latent size.} + +\item{patch_size}{Integer spatial patch size. Default 2.} + +\item{f_patch_size}{Integer temporal patch size. Default 1.} + +\item{out_channels}{Integer number of latent channels. Default 16.} +} +\value{ +Tensor of shape [C, F, H, W]. +} +\description{ +Takes the first F/pF * H/p * W/p tokens (the image span of the +unified sequence) and reassembles [C, F, H, W], matching unpatchify. +} diff --git a/tools/compare_translation.R b/tools/compare_translation.R index ceab7a2..7ffda61 100644 --- a/tools/compare_translation.R +++ b/tools/compare_translation.R @@ -208,7 +208,22 @@ report_pair("flux2 klein pipeline", py_scopes(pipe2, c("Flux2KleinPipeline", "compute_empirical_mu")), c("R/txt2img_flux2.R", "R/vae_flux2.R", "R/rope_flux2.R")) -# 5. Qwen3 encoder (modular reference; requires tools/cache/modular_qwen3.py) +# 5. Z-Image transformer stack + pipeline +tfz <- parse_file(file.path(DIFFUSERS, "models/transformers/transformer_z_image.py"), + ts_language_python()) +report_pair("zimage transformer", + py_scopes(tfz, c("TimestepEmbedder", "ZSingleStreamAttnProcessor", + "FeedForward", "ZImageTransformerBlock", "FinalLayer", + "RopeEmbedder", "ZImageTransformer2DModel")), + c("R/dit_zimage_modules.R", "R/dit_zimage.R", "R/rope_zimage.R")) + +pipez <- parse_file(file.path(DIFFUSERS, "pipelines/z_image/pipeline_z_image.py"), + ts_language_python()) +report_pair("zimage pipeline", + py_scopes(pipez, c("ZImagePipeline", "get_default_z_image_sigmas")), + c("R/txt2img_zimage.R")) + +# 6. Qwen3 encoder (modular reference; requires tools/cache/modular_qwen3.py) q3_path <- "tools/cache/modular_qwen3.py" if (file.exists(q3_path)) { q3 <- parse_file(q3_path, ts_language_python()) diff --git a/tools/gen_fixtures_zimage.py b/tools/gen_fixtures_zimage.py new file mode 100644 index 0000000..48985f8 --- /dev/null +++ b/tools/gen_fixtures_zimage.py @@ -0,0 +1,132 @@ +# Generate Z-Image phase-1 parity fixtures for the R port: 3-axis RoPE +# (theta 256, f32-angle cast), coordinate-grid position ids with the +# SEQ_MULTI_OF padding scheme, patchify/unpatchify, the Z-Image timestep +# sinusoid, and the static shift-3.0 sigma schedule. +# +# Runs the diffusers reference (Apache-2.0) on small fixed inputs. Run +# via tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.transformers.transformer_z_image import ( # noqa: E402 + RopeEmbedder, + TimestepEmbedder, + ZImageTransformer2DModel, +) +from diffusers.pipelines.z_image.pipeline_z_image import ( # noqa: E402 + get_default_z_image_sigmas, +) +from diffusers.schedulers import FlowMatchEulerDiscreteScheduler # noqa: E402 + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(53) +fx = {} + +# --- position ids via patchify_and_embed (tiny model, real geometry) ----------- +# cap len 37 pads to 64; image 6x10 = 60 tokens pads to 64. Asymmetric H/W +# so an axis swap fails. +tiny = ZImageTransformer2DModel( + in_channels=16, + dim=32, + n_layers=1, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + cap_feat_dim=24, + axes_dims=[4, 6, 6], + axes_lens=[64, 32, 32], +) + +CAP_LEN = 37 +IMG = torch.randn(16, 1, 12, 20) +CAP = torch.randn(CAP_LEN, 24) + +( + all_img_out, + all_cap_out, + all_img_size, + all_img_pos_ids, + all_cap_pos_ids, + all_img_pad_mask, + all_cap_pad_mask, +) = tiny.patchify_and_embed([IMG], [CAP], patch_size=2, f_patch_size=1) + +fx["img"] = IMG +fx["cap"] = CAP +fx["img_tokens_padded"] = all_img_out[0] # [64, 64] pads = last patch +fx["cap_feats_padded"] = all_cap_out[0] # [64, 24] pads = last row +fx["img_pos_ids"] = all_img_pos_ids[0].float() # [64, 3], pads (0,0,0) +# cap pos ids carry a dead (0,0,0) tail that _prepare_sequence truncates; +# keep only the effective first cap_padded rows (the 1..64 ramp) +fx["cap_pos_ids"] = all_cap_pos_ids[0][:64].float() # [64, 3] +fx["img_pad_mask"] = all_img_pad_mask[0].float() +fx["cap_pad_mask"] = all_cap_pad_mask[0].float() + +# --- unpatchify: slices [:ori_len] off the unified sequence --------------------- +uni_tokens = torch.randn(80, 2 * 2 * 1 * 16) # 64 img tokens + cap tail +unpat = tiny.unpatchify([uni_tokens.clone()], [(1, 12, 20)], patch_size=2, f_patch_size=1)[0] +fx["unpat_tokens"] = uni_tokens +fx["unpat_out"] = unpat # [16, 1, 12, 20] + +# --- RopeEmbedder at the real Turbo config -------------------------------------- +# Positions include the axis maxima to expose the f32-angle-cast rounding. +rope = RopeEmbedder(theta=256.0, axes_dims=[32, 48, 48], axes_lens=[1536, 512, 512]) +ids = torch.tensor( + [ + [1, 0, 0], + [2, 0, 0], + [64, 0, 0], + [65, 0, 0], + [65, 3, 7], + [65, 5, 9], + [130, 45, 60], + [1535, 511, 511], + [0, 0, 0], + ], + dtype=torch.int32, +) +freqs_cis = rope(ids) # [9, 64] complex64 +fx["rope_ids"] = ids.float() +fx["rope_cos"] = freqs_cis.real.repeat_interleave(2, dim=-1) # [9, 128] +fx["rope_sin"] = freqs_cis.imag.repeat_interleave(2, dim=-1) + + +# --- RoPE application (processor convention: x [B, S, H, D]) --------------------- +def apply_rotary_emb(x_in, freqs_cis): + x = torch.view_as_complex(x_in.float().reshape(*x_in.shape[:-1], -1, 2)) + freqs_cis = freqs_cis.unsqueeze(2) + x_out = torch.view_as_real(x * freqs_cis).flatten(3) + return x_out.type_as(x_in) + + +rope_x = torch.randn(1, 9, 3, 128) +fx["rope_x"] = rope_x +fx["rope_out"] = apply_rotary_emb(rope_x, freqs_cis.unsqueeze(0)) + +# --- timestep sinusoid ------------------------------------------------------------ +# Model input is t * t_scale with pipeline t = (1000 - t_sched)/1000 in [0, 1]. +t_in = torch.tensor([0.0, 125.0, 437.5, 875.0, 1000.0], dtype=torch.float32) +fx["t_emb_in"] = t_in +fx["t_emb_out"] = TimestepEmbedder.timestep_embedding(t_in, 256) # [5, 256] + +# --- scheduler: static shift 3.0 on linspace(1, 1/N, N) -------------------------- +for n in (4, 8): + sched = FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, shift=3.0, use_dynamic_shifting=False + ) + sched.set_timesteps(sigmas=get_default_z_image_sigmas(n)) + fx[f"sched_sigmas_{n}"] = sched.sigmas.float() # [n + 1] + fx[f"sched_timesteps_{n}"] = sched.timesteps.float() # [n] + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "rope_zimage.safetensors"), + metadata={"purpose": "diffuseR Z-Image test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/rope_zimage.safetensors") diff --git a/tools/gen_fixtures_zimage_dit.py b/tools/gen_fixtures_zimage_dit.py new file mode 100644 index 0000000..21fe7af --- /dev/null +++ b/tools/gen_fixtures_zimage_dit.py @@ -0,0 +1,127 @@ +# Generate Z-Image phase-2 parity fixtures: transformer block (modulated +# and unmodulated), final layer, timestep embedder, and cap embedder, +# all tiny random-init. +# +# Runs the diffusers reference (Apache-2.0) on small fixed inputs. Run +# via tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import sys + +import torch +import torch.nn as nn +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.normalization import RMSNorm # noqa: E402 +from diffusers.models.transformers.transformer_z_image import ( # noqa: E402 + FinalLayer, + RopeEmbedder, + TimestepEmbedder, + ZImageTransformerBlock, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(59) +fx = {} + +DIM, HEADS, HEAD_DIM = 32, 2, 16 +SEQ = 24 + +# Shared inputs: token positions mixing a cap ramp and an image grid +rope = RopeEmbedder(theta=256.0, axes_dims=[4, 6, 6], axes_lens=[64, 32, 32]) +ids = torch.cat( + [ + torch.tensor([[i + 1, 0, 0] for i in range(8)], dtype=torch.int32), + torch.stack( + torch.meshgrid( + torch.arange(9, 10, dtype=torch.int32), + torch.arange(4, dtype=torch.int32), + torch.arange(4, dtype=torch.int32), + indexing="ij", + ), + dim=-1, + ).flatten(0, 2), + ] +) +assert ids.shape == (SEQ, 3) +freqs_cis = rope(ids) # [24, 8] complex64 + +x = torch.randn(1, SEQ, DIM) +adaln = torch.randn(1, DIM) + +fx["ids"] = ids.float() +fx["freqs_cos"] = freqs_cis.real.repeat_interleave(2, dim=-1) # [24, 16] +fx["freqs_sin"] = freqs_cis.imag.repeat_interleave(2, dim=-1) +fx["x"] = x +fx["adaln"] = adaln + +# --- modulated block (noise refiner / main trunk) -------------------------------- +torch.manual_seed(61) +blk = ZImageTransformerBlock(0, DIM, HEADS, HEADS, norm_eps=1e-5, qk_norm=True, + modulation=True) +with torch.no_grad(): + for p in blk.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + out = blk(x, None, freqs_cis.unsqueeze(0), adaln) +for k, v in blk.state_dict().items(): + fx[f"mod.{k}"] = v +fx["mod_out"] = out + +# --- unmodulated block (context refiner) ------------------------------------------ +torch.manual_seed(67) +ublk = ZImageTransformerBlock(0, DIM, HEADS, HEADS, norm_eps=1e-5, qk_norm=True, + modulation=False) +with torch.no_grad(): + for p in ublk.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + uout = ublk(x, None, freqs_cis.unsqueeze(0)) +for k, v in ublk.state_dict().items(): + fx[f"unmod.{k}"] = v +fx["unmod_out"] = uout + +# --- final layer ------------------------------------------------------------------- +torch.manual_seed(71) +fin = FinalLayer(DIM, 2 * 2 * 1 * 16) +with torch.no_grad(): + for p in fin.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + fout = fin(x, c=adaln) +for k, v in fin.state_dict().items(): + fx[f"final.{k}"] = v +fx["final_out"] = fout + +# --- timestep embedder -------------------------------------------------------------- +torch.manual_seed(73) +temb = TimestepEmbedder(min(DIM, 256), mid_size=48) +with torch.no_grad(): + for p in temb.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + t_in = torch.tensor([0.0, 437.5, 1000.0]) + t_out = temb(t_in) +for k, v in temb.state_dict().items(): + fx[f"temb.{k}"] = v +fx["temb_in"] = t_in +fx["temb_out"] = t_out + +# --- cap embedder (RMSNorm + Linear, Sequential keys 0/1) ---------------------------- +torch.manual_seed(79) +CAP_DIM = 24 +cap_embedder = nn.Sequential(RMSNorm(CAP_DIM, eps=1e-5), nn.Linear(CAP_DIM, DIM, bias=True)) +with torch.no_grad(): + for p in cap_embedder.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + cap_in = torch.randn(5, CAP_DIM) + cap_out = cap_embedder(cap_in) +for k, v in cap_embedder.state_dict().items(): + fx[f"cap.{k}"] = v +fx["cap_in"] = cap_in +fx["cap_out"] = cap_out + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "dit_zimage.safetensors"), + metadata={"purpose": "diffuseR Z-Image test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/dit_zimage.safetensors") diff --git a/tools/gen_fixtures_zimage_model.py b/tools/gen_fixtures_zimage_model.py new file mode 100644 index 0000000..64d719a --- /dev/null +++ b/tools/gen_fixtures_zimage_model.py @@ -0,0 +1,64 @@ +# Generate the Z-Image phase-3 fixture: a tiny random-init +# ZImageTransformer2DModel full forward (state dict + input + output), +# with both caption and image needing pad tokens, plus a sharded +# save_pretrained checkpoint for loader/quantizer tests. +# +# Runs the diffusers reference (Apache-2.0) on small fixed inputs. Run +# via tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import sys + +import torch +from safetensors.torch import save_file + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src")) + +from diffusers.models.transformers.transformer_z_image import ( # noqa: E402 + ZImageTransformer2DModel, +) + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(83) +fx = {} + +model = ZImageTransformer2DModel( + in_channels=4, + dim=48, + n_layers=2, + n_refiner_layers=1, + n_heads=2, + n_kv_heads=2, + cap_feat_dim=24, + axes_dims=[8, 8, 8], + axes_lens=[128, 32, 32], +) +with torch.no_grad(): + for p in model.parameters(): + p.copy_(torch.randn_like(p) * 0.05) + +# 12x20 latent -> 6x10 = 60 tokens, pads to 64; caption 37 pads to 64 +x = torch.randn(4, 1, 12, 20) +cap = torch.randn(37, 24) +t = torch.tensor([0.4375]) + +with torch.no_grad(): + out = model([x], t, [cap], return_dict=False)[0][0] + +for k, v in model.state_dict().items(): + fx[f"sd.{k}"] = v +fx["x"] = x +fx["cap"] = cap +fx["t"] = t +fx["out"] = out # [4, 1, 12, 20] + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "zimage_model.safetensors"), + metadata={"purpose": "diffuseR Z-Image test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/zimage_model.safetensors") + +CKPT_DIR = os.path.join(OUT_DIR, "zimage_tiny_ckpt") +model.save_pretrained(CKPT_DIR, max_shard_size="120KB") +print(f"wrote sharded checkpoint to {CKPT_DIR}") diff --git a/tools/gen_fixtures_zimage_qwen.py b/tools/gen_fixtures_zimage_qwen.py new file mode 100644 index 0000000..36a32eb --- /dev/null +++ b/tools/gen_fixtures_zimage_qwen.py @@ -0,0 +1,61 @@ +# Generate the Z-Image phase-4 fixture: a tiny random-init Qwen3Model +# pinning the hidden_states[-2] index convention and the +# mask-slice-to-variable-length caption used by the Z-Image pipeline. +# +# Run via tools/gen_fixtures.sh; never executed at package test/run time. + +import os +import sys + +import torch +from safetensors.torch import save_file +from transformers import Qwen3Config, Qwen3Model + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(89) +fx = {} + +N_LAYERS = 4 +config = Qwen3Config( + vocab_size=128, + hidden_size=32, + intermediate_size=64, + num_hidden_layers=N_LAYERS, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + rope_theta=1e6, + rms_norm_eps=1e-6, + tie_word_embeddings=True, +) +model = Qwen3Model(config) +with torch.no_grad(): + for p in model.parameters(): + p.copy_(torch.randn_like(p) * 0.05) +model.eval() + +# 9 real tokens + 3 right pads +input_ids = torch.tensor([[5, 17, 99, 3, 42, 8, 120, 64, 7, 0, 0, 0]]) +attention_mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0]]) + +with torch.no_grad(): + out = model(input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True) + +# hidden_states has N_LAYERS + 1 entries; [-2] = after layer N_LAYERS - 1 +penult = out.hidden_states[-2] +assert torch.equal(penult, out.hidden_states[N_LAYERS - 1]) + +for k, v in model.state_dict().items(): + fx[f"sd.{k}"] = v +fx["input_ids"] = input_ids.float() +fx["attention_mask"] = attention_mask.float() +fx["penult"] = penult # [1, 12, 32] +fx["penult_sliced"] = penult[0][attention_mask[0].bool()] # [9, 32] + +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "zimage_qwen.safetensors"), + metadata={"purpose": "diffuseR Z-Image test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/zimage_qwen.safetensors") diff --git a/tools/gen_zimage_qwen_template_cases.py b/tools/gen_zimage_qwen_template_cases.py new file mode 100644 index 0000000..0ace1bf --- /dev/null +++ b/tools/gen_zimage_qwen_template_cases.py @@ -0,0 +1,61 @@ +# Generate Z-Image chat-template parity cases: the Qwen3 template with +# enable_thinking=True (no think block), rendered by the SHIPPED +# Tongyi-MAI/Z-Image-Turbo tokenizer_config and padded like the +# pipeline. The vocab/merges/tokenizer.json are byte-identical to +# FLUX.2-klein's (verified by blob oid), so raw BPE parity is already +# covered by qwen_tokenizer_cases.json; only the template render is +# pinned here. +# +# Writes inst/tinytest/fixtures/zimage_template_cases.json (checked in). +# +# Run: +# uv run --no-project --with transformers --with torch \ +# --index https://download.pytorch.org/whl/cpu \ +# --index-strategy unsafe-best-match python tools/gen_zimage_qwen_template_cases.py + +import json +import os + +from transformers import AutoTokenizer + +ROOT = os.path.join(os.path.dirname(__file__), "..") +FIXTURE = os.path.join(ROOT, "inst", "tinytest", "fixtures", "zimage_template_cases.json") + +tok = AutoTokenizer.from_pretrained("Tongyi-MAI/Z-Image-Turbo", subfolder="tokenizer") + +PROMPTS = [ + "a photo of a cat", + "emoji 🦊 and 中文字符 mixed in", + "An astronaut riding a horse on Mars, photorealistic", + "一幅为名为“造相「Z-IMAGE-TURBO」”的项目设计的创意海报。", + "", +] + +templated = [] +for text in PROMPTS: + messages = [{"role": "user", "content": text}] + rendered = tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, + enable_thinking=True, + ) + enc = tok(rendered, padding="max_length", truncation=True, max_length=64) + templated.append({ + "text": text, + "rendered": rendered, + "max_length": 64, + "ids": enc["input_ids"], + "mask": enc["attention_mask"], + }) + +meta = { + "pad_token": tok.pad_token, + "pad_token_id": tok.pad_token_id, + "padding_side": tok.padding_side, +} + +with open(FIXTURE, "w", encoding="utf-8") as f: + json.dump({"templated": templated, "meta": meta}, + f, ensure_ascii=False, indent=1) +print(f"wrote {len(templated)} templated cases to {FIXTURE}") +print("meta:", meta) +print("rendered example:", json.dumps(templated[0]["rendered"]))