diff --git a/DESCRIPTION b/DESCRIPTION index 923d34c..aa5c3d3 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.13 +Version: 0.2.2.14 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 3bb9f58..439c73e 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -18,9 +18,11 @@ export(download_sd21) export(download_sdxl) export(download_zimage_turbo) export(encode_bpe) +export(encode_gemma4) export(encode_qwen) export(encode_unigram) export(encode_with_gemma3) +export(encode_with_gemma4) export(encode_with_qwen3) export(encode_with_t5) export(filename_from_prompt) @@ -70,6 +72,9 @@ export(gemma3_encode_batch) export(gemma3_quantize_nf4) export(gemma3_text_model) export(gemma3_tokenizer) +export(gemma4_quantize_nf4) +export(gemma4_text_model) +export(gemma4_tokenizer) export(img2img) export(is_blackwell_gpu) export(latents_to_video) @@ -78,6 +83,8 @@ export(load_decoder_weights) export(load_flux2_vae_decoder) export(load_gemma3_nf4) export(load_gemma3_text_encoder) +export(load_gemma4_nf4) +export(load_gemma4_text_encoder) export(load_model_component) export(load_pipeline) export(load_qwen3_text_encoder) @@ -110,6 +117,7 @@ export(ltx23_census) export(ltx23_connector_transformer_1d) export(ltx23_denormalize_latents) export(ltx23_distilled_sigmas) +export(ltx23_dit_config_args) export(ltx23_downsample1d) export(ltx23_encode_audio) export(ltx23_encode_video_frames) @@ -172,6 +180,7 @@ export(ltx23_video_vae) export(ltx23_vocoder) export(ltx23_vocoder_resblock) export(ltx23_vocoder_with_bwe) +export(ltx25_open_checkpoint) export(models2devices) export(offload_to_cpu) export(post_quant_conv) @@ -238,6 +247,7 @@ export(zimage_unpatchify) S3method(print,bpe_tokenizer) S3method(print,diffuseR_resident) +S3method(print,gemma4_tokenizer) S3method(print,ltx23_checkpoint) S3method(print,qwen_tokenizer) S3method(print,unigram_tokenizer) diff --git a/NEWS.md b/NEWS.md index a9d6d3f..f1bca17 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,28 @@ +# diffuseR 0.2.2.14 + +* **LTX-2.5 (distilled) text-to-video.** The 2.3 pipeline runs the 2.5 + checkpoints: + - Transformer flags come from the embedded config (`ff_bias = FALSE`, + keyframe position embedding); `ltx23_open_checkpoint()` accepts + model version 2.5. + - `txt2vid_ltx2(sampler = "auto")` uses the ancestral Euler sampler + for 2.5 (stage noise from `seed + 10000` / `seed + 20000`), Euler + for 2.3. + - New Gemma 4 text encoder: `gemma4_text_model()`, + `load_gemma4_text_encoder()`, `encode_with_gemma4()`, an exact + ranked-BPE `gemma4_tokenizer()`, and `gemma4_quantize_nf4()` / + `load_gemma4_nf4()`. `txt2vid_ltx2()` picks Gemma 4 by tokenizer + class. + - `ltx25_open_checkpoint()` presents the split 2.5 files (transformer, + encoder file with the text projection, conv VAE, audio VAE + + vocoder) as one checkpoint; `ltx23_load_pipeline()` takes it, and + `ltx23_quantize_nf4()` writes it as one self-contained NF4 artifact. + - `resident_load("ltx")` loads a Gemma 4 encoder artifact. + - The diffusion-decoder VAE, duration head and temporal upsampler are + not supported yet. +* The LTX fp8/NF4 and Gemma 4 quantizers collect garbage per tensor, so + quantizing no longer outgrows host RAM. + # diffuseR 0.2.2.13 * **`zimage_load_pipeline()` takes `revision`**, as diff --git a/R/checkpoint_ltx23.R b/R/checkpoint_ltx23.R index 59d2619..3c14762 100644 --- a/R/checkpoint_ltx23.R +++ b/R/checkpoint_ltx23.R @@ -18,8 +18,9 @@ NULL #' configuration. #' #' @param path Path to the checkpoint .safetensors file. -#' @param require_version Character. Required \code{model_version} prefix -#' (default "2.3"). Set to NULL to skip the check. +#' @param require_version Character vector. Accepted \code{model_version} +#' prefixes (default \code{c("2.3", "2.5")}, the 22B generations this +#' module tree loads). Set to NULL to skip the check. #' #' @return An object of class \code{ltx23_checkpoint}: a list with #' \code{handle} (safetensors reader), \code{keys}, \code{version}, @@ -32,7 +33,7 @@ NULL #' ckpt <- ltx23_open_checkpoint("ltx-2.3-22b-distilled-1.1.safetensors") #' str(ltx23_split_keys(ckpt$keys), max.level = 1) #' } -ltx23_open_checkpoint <- function(path, require_version = "2.3") { +ltx23_open_checkpoint <- function(path, require_version = c("2.3", "2.5")) { if (!requireNamespace("safetensors", quietly = TRUE)) { stop("The safetensors package is required to read LTX checkpoints.") } @@ -51,14 +52,15 @@ ltx23_open_checkpoint <- function(path, require_version = "2.3") { } if (!is.null(require_version)) { + wanted <- paste0(require_version, "*", collapse = " or ") if (is.null(version)) { stop("Checkpoint has no model_version metadata; expected an LTX ", - require_version, " single-file checkpoint: ", path) + wanted, " single-file checkpoint: ", path) } - if (!startsWith(version, require_version)) { + if (!any(startsWith(version, require_version))) { stop( - "Checkpoint model_version is '", version, "' but '", - require_version, "*' is required: ", path + "Checkpoint model_version is '", version, "' but ", + wanted, " is required: ", path ) } } diff --git a/R/checkpoint_ltx25.R b/R/checkpoint_ltx25.R new file mode 100644 index 0000000..bf8995d --- /dev/null +++ b/R/checkpoint_ltx25.R @@ -0,0 +1,102 @@ +#' Open the LTX-2.5 component files as one checkpoint +#' +#' LTX-2.3 shipped every component in one file; LTX-2.5 splits them: the +#' transformer (with the text connectors), the Gemma 4 text encoder file +#' (which carries the connectors' \code{text_embedding_projection}), the +#' video VAE, and the audio VAE plus vocoder. This presents them as a +#' single \code{ltx23_checkpoint} under the 2.3 key names (the video VAE +#' keys gain their \code{vae.} prefix), so \code{\link{ltx23_load_pipeline}} +#' and its component loaders work unchanged. +#' +#' @param transformer The 2.5 transformer file +#' (\code{ltx-2.5-22b-distilled-transformer-bf16.safetensors}) or an +#' fp8/NF4 artifact directory quantized from it. +#' @param text_encoder The 2.5 Gemma 4 text encoder file, or its NF4 +#' artifact directory (\code{\link{gemma4_quantize_nf4}} keeps the +#' projection); only \code{text_embedding_projection.*} is used here. +#' @param video_vae The 2.5 video VAE file. Use the convolutional decoder +#' (\code{ltx-2.5-video-vae-conv-bf16.safetensors}); the default +#' diffusion decoder is not supported. +#' @param audio_vae Optional 2.5 audio VAE + vocoder file +#' (\code{ltx-2.5-audio-vae-bf16.safetensors}). +#' +#' @return An \code{ltx23_checkpoint} whose \code{handle$get_tensor()} +#' reads each key from its source file. +#' +#' @export +ltx25_open_checkpoint <- function(transformer, text_encoder, video_vae, + audio_vae = NULL) { + open_any <- function(path, require_version) { + if (file.exists(file.path(path, "manifest.json"))) { + ltx23_open_fp8_checkpoint(path) + } else { + ltx23_open_checkpoint(path, require_version = require_version) + } + } + dit <- open_any(transformer, "2.5") + if (!startsWith(dit$version %||% "", "2.5")) { + stop("Transformer is model_version '", dit$version %||% "none", + "', expected 2.5*: ", transformer) + } + te <- open_any(text_encoder, NULL) + vae <- ltx23_open_checkpoint(video_vae, require_version = "2.5") + if (any(grepl("(^|\\.)transformer_blocks\\.", vae$keys)) || + !any(grepl("^(vae\\.)?decoder\\.up_blocks\\.", vae$keys))) { + stop("Not a convolutional video VAE file (the diffusion decoder is ", + "not supported): ", video_vae) + } + sources <- list(dit = dit, te = te, vae = vae) + if (!is.null(audio_vae)) { + sources$audio <- ltx23_open_checkpoint(audio_vae, require_version = "2.5") + } + + # canonical key -> (source, original key) + proj <- grep("^text_embedding_projection\\.", te$keys, value = TRUE) + if (!length(proj)) { + stop("No text_embedding_projection.* in ", text_encoder) + } + vae_keys <- vae$keys + routes <- list( + dit = list(orig = dit$keys, canon = dit$keys), + te = list(orig = proj, canon = proj), + vae = list(orig = vae_keys, + canon = ifelse(startsWith(vae_keys, "vae."), vae_keys, + paste0("vae.", vae_keys))) + ) + if (!is.null(sources$audio)) { + ak <- grep("^(audio_vae|vocoder)\\.", sources$audio$keys, value = TRUE) + routes$audio <- list(orig = ak, canon = ak) + } + canon <- unlist(lapply(routes, `[[`, "canon"), use.names = FALSE) + orig <- unlist(lapply(routes, `[[`, "orig"), use.names = FALSE) + src <- rep(names(routes), vapply(routes, function(r) length(r$canon), 1L)) + dup <- canon[duplicated(canon)] + if (length(dup)) { + stop("Key present in more than one LTX-2.5 file: ", dup[1]) + } + + handle <- list(get_tensor = function(key) { + i <- match(key, canon) + if (is.na(i)) { + stop("Key not found in the LTX-2.5 files: ", key) + } + sources[[src[i]]]$handle$get_tensor(orig[i]) + }) + + vcfg <- vae$config + acfg <- sources$audio$config + config <- list(transformer = dit$config$transformer, + scheduler = dit$config$scheduler, + vae = vcfg$vae %||% vcfg, + audio_vae = acfg$audio_vae, vocoder = acfg$vocoder) + structure( + list(handle = handle, keys = canon, version = dit$version, + config = config, format = dit$format, path = transformer, + sources = vapply(list(transformer = transformer, + text_encoder = text_encoder, + video_vae = video_vae, + audio_vae = audio_vae %||% NA_character_), + path.expand, "")), + class = "ltx23_checkpoint" + ) +} diff --git a/R/dit_ltx23.R b/R/dit_ltx23.R index 336e0c3..fad921d 100644 --- a/R/dit_ltx23.R +++ b/R/dit_ltx23.R @@ -36,6 +36,15 @@ NULL #' @param rope_type "split" (LTX-2.3) or "interleaved". #' @param gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn #' LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints). +#' @param ff_bias,audio_ff_bias Logicals. Feed-forward biases per stream +#' (\code{ff_bias = FALSE} for LTX-2.5). +#' @param use_keyframes_abs_pos_embedding Logical. Carry the LTX-2.5 +#' learned keyframe marker, added after \code{proj_in} to the tokens +#' flagged by the forward's \code{video_keyframes_mask} +#' (\code{[B, S, 1]}, non-zero on keyframe tokens). +#' +#' @seealso \code{\link{ltx23_dit_config_args}} to derive these from a +#' checkpoint's embedded config. #' #' @return Module whose forward(hidden_states, ...) returns #' \code{list(sample, audio_sample)}: the predicted velocity for the @@ -80,7 +89,10 @@ ltx23_transformer <- torch::nn_module( timestep_scale_multiplier = 1000, cross_attn_timestep_scale_multiplier = 1000, rope_type = "split", - perturbed_attn = TRUE + perturbed_attn = TRUE, + ff_bias = TRUE, + audio_ff_bias = TRUE, + use_keyframes_abs_pos_embedding = FALSE ) { inner_dim <- num_attention_heads * attention_head_dim audio_inner_dim <- audio_num_attention_heads * audio_attention_head_dim @@ -94,6 +106,13 @@ ltx23_transformer <- torch::nn_module( self$proj_in <- linear_noinit(in_channels, inner_dim) self$audio_proj_in <- linear_noinit(audio_in_channels, audio_inner_dim) + # LTX-2.5 keyframe marker (zero-initialized in the reference) + if (use_keyframes_abs_pos_embedding) { + self$keyframes_abs_pos_embedding <- torch::nn_parameter( + torch::torch_zeros(1L, inner_dim) + ) + } + # 2. Global timestep modulation (9 params each with cross-attn mod) video_mod_params <- if (cross_attn_mod) 9L else 6L audio_mod_params <- if (audio_cross_attn_mod) 9L else 6L @@ -175,7 +194,9 @@ ltx23_transformer <- torch::nn_module( eps = norm_eps, elementwise_affine = FALSE, rope_type = rope_type, - perturbed_attn = perturbed_attn + perturbed_attn = perturbed_attn, + ff_bias = ff_bias, + audio_ff_bias = audio_ff_bias ) })) @@ -209,7 +230,8 @@ ltx23_transformer <- torch::nn_module( perturbation_mask = NULL, use_cross_timestep = FALSE, video_self_attention_mask = NULL, - cond_token_index = NULL + cond_token_index = NULL, + video_keyframes_mask = NULL ) { audio_timestep <- audio_timestep %||% timestep audio_sigma <- audio_sigma %||% sigma @@ -262,6 +284,15 @@ ltx23_transformer <- torch::nn_module( hidden_states <- self$proj_in(hidden_states) audio_hidden_states <- self$audio_proj_in(audio_hidden_states) + # 2.1. LTX-2.5: mark tokens whose latent frame encodes a single pixel + # frame (causal first frame, generated keyframe slots) + if (!is.null(self$keyframes_abs_pos_embedding) && + !is.null(video_keyframes_mask)) { + marker <- video_keyframes_mask$gt(0)$to(dtype = hidden_states$dtype) + hidden_states <- hidden_states + marker * + self$keyframes_abs_pos_embedding$to(dtype = hidden_states$dtype) + } + # 3. Timestep embeddings and global modulation parameters gate_scale_factor <- self$cross_attn_timestep_scale_multiplier / self$timestep_scale_multiplier @@ -433,6 +464,55 @@ ltx23_transformer <- torch::nn_module( } ) +#' Transformer arguments from a checkpoint's embedded config +#' +#' Reads the version-dependent flags out of the \code{transformer} entry +#' of an LTX single-file checkpoint's embedded config. Every other +#' dimension is shared by the 2.3 and 2.5 22B checkpoints and stays at +#' the \code{\link{ltx23_transformer}} defaults. Keys absent from the +#' config (any 2.3 checkpoint) are left out, so the defaults apply. +#' +#' @param config The parsed embedded config (\code{ckpt$config}), or NULL. +#' +#' @return Named list of \code{\link{ltx23_transformer}} arguments. +#' +#' @export +ltx23_dit_config_args <- function(config) { + tc <- config$transformer + keys <- c("ff_bias", "audio_ff_bias", "use_keyframes_abs_pos_embedding") + args <- list() + for (k in keys) { + if (!is.null(tc[[k]])) { + args[[k]] <- isTRUE(tc[[k]]) + } + } + args +} + +# Load the unquantized transformer from a single-file checkpoint: a +# skeleton built from the embedded config (explicit ... wins, as in the +# fp8/NF4 loaders), streamed in at `dtype`, refused unless every key +# lands and every parameter is filled. +.ltx23_load_transformer <- function(ckpt, dtype, device = "cpu", + verbose = TRUE, ...) { + args <- utils::modifyList(ltx23_dit_config_args(ckpt$config), list(...)) + model <- do.call(.construct_skeleton, + c(list(ltx23_transformer, dtype = dtype), args)) + keys <- ltx23_split_keys(ckpt$keys)$dit + if (verbose) { + message("Loading dit (", length(keys), " tensors)...") + } + res <- ltx23_load_group(ckpt, keys, model, map_key = ltx23_map_dit_key, + verbose = verbose) + if (length(res$unmapped) || length(res$unfilled)) { + stop("dit: incomplete load (", length(res$unmapped), " unmapped, ", + length(res$unfilled), " unfilled)") + } + model$to(device = device) + model$eval() + model +} + #' Map an official DiT checkpoint key to the R module name #' #' Applies the official-to-diffusers renames for the LTX-2.3 transformer diff --git a/R/dit_ltx23_modules.R b/R/dit_ltx23_modules.R index 93288c6..35dff3e 100644 --- a/R/dit_ltx23_modules.R +++ b/R/dit_ltx23_modules.R @@ -491,6 +491,8 @@ ltx23_attention <- torch::nn_module( #' #' @param dim Integer. Input/output dimension. #' @param mult Integer. Hidden dimension multiplier. +#' @param bias Logical. Both linears carry a bias (FALSE for the LTX-2.5 +#' video stream). #' #' @return Module whose forward(x) returns the projected states passed #' through a tanh-approximated GELU, a tensor of the same shape as @@ -499,12 +501,12 @@ ltx23_attention <- torch::nn_module( #' @export ltx23_feed_forward <- torch::nn_module( "ltx23_feed_forward", - initialize = function(dim, mult = 4L) { + initialize = function(dim, mult = 4L, bias = TRUE) { inner_dim <- as.integer(dim * mult) gelu_proj <- torch::nn_module( "ltx23_gelu", initialize = function(dim_in, dim_out) { - self$proj <- linear_noinit(dim_in, dim_out) + self$proj <- linear_noinit(dim_in, dim_out, bias = bias) }, forward = function(x) { h <- self$proj(x) @@ -517,7 +519,7 @@ ltx23_feed_forward <- torch::nn_module( self$net <- torch::nn_module_list(list( gelu_proj(dim, inner_dim), torch::nn_identity(), # dropout slot in the reference; inference no-op - linear_noinit(inner_dim, dim) + linear_noinit(inner_dim, dim, bias = bias) )) }, forward = function(x) { @@ -554,6 +556,8 @@ ltx23_get_mod_params <- function(scale_shift_table, temb, batch_size) { #' @param elementwise_affine Logical. Block norms carry weights (FALSE for LTX). #' @param rope_type "split" or "interleaved". #' @param perturbed_attn Logical. Enable the STG perturbation arguments. +#' @param ff_bias,audio_ff_bias Logicals. Feed-forward biases per stream +#' (LTX-2.5 drops the video ones). #' #' @return Module whose forward(hidden_states, ...) returns #' \code{list(hidden_states, audio_hidden_states)}, the video and audio @@ -579,7 +583,9 @@ ltx23_transformer_block <- torch::nn_module( eps = 1e-6, elementwise_affine = FALSE, rope_type = "split", - perturbed_attn = TRUE + perturbed_attn = TRUE, + ff_bias = TRUE, + audio_ff_bias = TRUE ) { self$perturbed_attn <- perturbed_attn @@ -627,9 +633,9 @@ ltx23_transformer_block <- torch::nn_module( # 4. Feed-forward self$norm3 <- ltx23_rms_norm(dim, eps = eps, elementwise_affine = elementwise_affine) - self$ff <- ltx23_feed_forward(dim) + self$ff <- ltx23_feed_forward(dim, bias = ff_bias) self$audio_norm3 <- ltx23_rms_norm(audio_dim, eps = eps, elementwise_affine = elementwise_affine) - self$audio_ff <- ltx23_feed_forward(audio_dim) + self$audio_ff <- ltx23_feed_forward(audio_dim, bias = audio_ff_bias) # 5. Per-block modulation tables self$video_cross_attn_adaln <- video_cross_attn_adaln diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index 4a7fece..aeb737e 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -193,11 +193,11 @@ ltx23_quantize_fp8 <- function(checkpoint_path, output_dir = NULL, if (shard_size >= shard_bytes) { flush_shard() } - if (i %% 200L == 0L) { - gc(verbose = FALSE) - if (verbose) { - message(sprintf(" quantizing %d/%d tensors", i, length(keys))) - } + # Per tensor, as in ltx23_quantize_nf4(): R's gc cannot see + # torch's native memory + gc(verbose = FALSE) + if (i %% 200L == 0L && verbose) { + message(sprintf(" quantizing %d/%d tensors", i, length(keys))) } } flush_shard() @@ -287,7 +287,8 @@ ltx23_open_fp8_checkpoint <- function(dir) { ltx23_load_transformer_fp8 <- function(ckpt, device = "cuda", pin = TRUE, verbose = TRUE, ...) { stopifnot(inherits(ckpt, "ltx23_checkpoint")) - model <- .construct_skeleton(ltx23_transformer, ...) + args <- utils::modifyList(ltx23_dit_config_args(ckpt$config), list(...)) + model <- do.call(.construct_skeleton, c(list(ltx23_transformer), args)) groups <- ltx23_split_keys(ckpt$keys) dit_keys <- groups$dit diff --git a/R/gemma4_text_encoder.R b/R/gemma4_text_encoder.R new file mode 100644 index 0000000..4f12654 --- /dev/null +++ b/R/gemma4_text_encoder.R @@ -0,0 +1,447 @@ +# Gemma 4 Text Encoder for LTX-2.5 +# +# Native R torch port of the text stack of HuggingFace transformers' +# Gemma4UnifiedTextModel (Apache-2.0, models/gemma4_unified). LTX-2.5 +# ships a fine-tuned 12B (gemma4-12b-ltx-v1) whose trunk matches the +# Gemma 3 12B of LTX-2.3 (48 layers, hidden 3840, 5 sliding : 1 full), +# with these differences: +# - RMSNorm scales by w (Gemma 3: 1 + w); values get a scale-free RMSNorm +# - attention scaling is 1.0 (the q/k norms carry the scale) +# - full-attention layers: head_dim 512, one KV head, keys reused as +# values (attention_k_eq_v), "proportional" RoPE rotating a quarter of +# the dims at theta 1e6 +# - every layer multiplies its output by a stored layer_scalar +# Vision/audio towers and the KV-sharing / per-layer-input / MoE options +# are off in this checkpoint and not ported. + +#' Gemma 4 RMS normalization +#' +#' @param dim Integer. Normalized dimension. +#' @param eps Numeric. Epsilon. +#' @param with_scale Logical. Carry a learned scale (FALSE for the value +#' norm). +#' +#' @return Module whose forward(x) normalizes over the last dimension in +#' float32 and multiplies by the scale, returning \code{x}'s dtype. +#' +#' @keywords internal +gemma4_rms_norm <- torch::nn_module( + "Gemma4RMSNorm", + initialize = function(dim, eps = 1e-6, with_scale = TRUE) { + self$eps <- eps + self$with_scale <- with_scale + if (with_scale) { + self$weight <- torch::nn_parameter(torch::torch_ones(dim)) + } +}, + forward = function(x) { + xf <- x$to(dtype = torch::torch_float32()) + out <- xf * xf$pow(2)$mean(dim = -1L, keepdim = TRUE)$add(self$eps)$pow(-0.5) + if (self$with_scale) { + out <- out * self$weight$to(dtype = torch::torch_float32()) + } + out$to(dtype = x$dtype) +} +) + +# Inverse RoPE frequencies for one layer type. "proportional" rotates +# the first floor(partial * head_dim / 2) frequency pairs, computed over +# the full head_dim, and leaves the rest unrotated (zero frequency). +.gemma4_inv_freq <- function(head_dim, params) { + base <- params$rope_theta + if (identical(params$rope_type %||% "default", "proportional")) { + angles <- as.integer(floor((params$partial_rotary_factor %||% 1) * + head_dim / 2)) + rotated <- 1 / base^(seq(0, 2 * angles - 1, by = 2) / head_dim) + c(rotated, rep(0, head_dim %/% 2L - angles)) + } else { + 1 / base^(seq(0, head_dim - 1, by = 2) / head_dim) + } +} + +# cos/sin [B, S, head_dim] for integer positions [B, S] +.gemma4_rope <- function(inv_freq, position_ids, dtype) { + inv <- torch::torch_tensor(inv_freq, dtype = torch::torch_float32(), + device = position_ids$device) + freqs <- position_ids$to(dtype = torch::torch_float32())$unsqueeze(-1L) * inv + emb <- torch::torch_cat(list(freqs, freqs), dim = -1L) + list(emb$cos()$to(dtype = dtype), emb$sin()$to(dtype = dtype)) +} + +#' Gemma 4 text attention +#' +#' @param config Model configuration (see \code{\link{load_gemma4_text_encoder}}). +#' @param layer_idx Integer. 0-based layer index. +#' +#' @return Module whose forward(hidden_states, attention_mask, cos, sin) +#' returns the attention output [B, S, hidden_size]. +#' +#' @keywords internal +gemma4_attention <- torch::nn_module( + "Gemma4Attention", + initialize = function(config, layer_idx) { + full <- identical(config$layer_types[[layer_idx + 1L]], "full_attention") + self$is_sliding <- !full + self$num_heads <- config$num_attention_heads + self$head_dim <- if (full) config$global_head_dim else config$head_dim + self$k_eq_v <- isTRUE(config$attention_k_eq_v) && full + self$num_kv_heads <- if (self$k_eq_v) { + config$num_global_key_value_heads %||% config$num_key_value_heads + } else { + config$num_key_value_heads + } + eps <- config$rms_norm_eps + h <- config$hidden_size + + self$q_proj <- linear_noinit(h, self$num_heads * self$head_dim, bias = FALSE) + self$k_proj <- linear_noinit(h, self$num_kv_heads * self$head_dim, bias = FALSE) + if (!self$k_eq_v) { + self$v_proj <- linear_noinit(h, self$num_kv_heads * self$head_dim, + bias = FALSE) + } + self$o_proj <- linear_noinit(self$num_heads * self$head_dim, h, bias = FALSE) + self$q_norm <- gemma4_rms_norm(self$head_dim, eps = eps) + self$k_norm <- gemma4_rms_norm(self$head_dim, eps = eps) + self$v_norm <- gemma4_rms_norm(self$head_dim, eps = eps, with_scale = FALSE) +}, + forward = function(hidden_states, attention_mask, cos, sin) { + b <- hidden_states$shape[1] + s <- hidden_states$shape[2] + heads <- function(x, n) x$view(c(b, s, n, self$head_dim)) + cos <- cos$unsqueeze(3L) + sin <- sin$unsqueeze(3L) + rope <- function(x) x * cos + rotate_half(x) * sin + + q <- rope(self$q_norm(heads(self$q_proj(hidden_states), self$num_heads))) + k_raw <- heads(self$k_proj(hidden_states), self$num_kv_heads) + # k == v layers take the values from the raw key projection + v_raw <- if (self$k_eq_v) { + k_raw + } else { + heads(self$v_proj(hidden_states), self$num_kv_heads) + } + k <- rope(self$k_norm(k_raw)) + v <- self$v_norm(v_raw) + + q <- q$transpose(2L, 3L) + k <- repeat_kv(k$transpose(2L, 3L), self$num_heads %/% self$num_kv_heads) + v <- repeat_kv(v$transpose(2L, 3L), self$num_heads %/% self$num_kv_heads) + out <- torch::torch_scaled_dot_product_attention(q, k, v, + attn_mask = attention_mask, + scale = 1.0) + out <- out$transpose(2L, 3L)$reshape(c(b, s, self$num_heads * self$head_dim)) + self$o_proj(out) +} +) + +#' Gemma 4 decoder layer +#' +#' @param config Model configuration. +#' @param layer_idx Integer. 0-based layer index. +#' +#' @return Module whose forward(hidden_states, attention_mask, cos, sin) +#' returns the layer output, scaled by the stored \code{layer_scalar}. +#' +#' @keywords internal +gemma4_decoder_layer <- torch::nn_module( + "Gemma4DecoderLayer", + initialize = function(config, layer_idx) { + eps <- config$rms_norm_eps + h <- config$hidden_size + self$self_attn <- gemma4_attention(config, layer_idx) + self$mlp <- gemma3_mlp(config) + self$input_layernorm <- gemma4_rms_norm(h, eps = eps) + self$post_attention_layernorm <- gemma4_rms_norm(h, eps = eps) + self$pre_feedforward_layernorm <- gemma4_rms_norm(h, eps = eps) + self$post_feedforward_layernorm <- gemma4_rms_norm(h, eps = eps) + self$layer_scalar <- torch::nn_buffer(torch::torch_ones(1L)) +}, + forward = function(hidden_states, attention_mask, cos, sin) { + residual <- hidden_states + x <- self$input_layernorm(hidden_states) + x <- self$self_attn(x, attention_mask, cos, sin) + hidden_states <- residual + self$post_attention_layernorm(x) + + residual <- hidden_states + x <- self$mlp(self$pre_feedforward_layernorm(hidden_states)) + hidden_states <- residual + self$post_feedforward_layernorm(x) + hidden_states * self$layer_scalar$to(dtype = hidden_states$dtype) +} +) + +#' Gemma 4 text model +#' +#' The text stack of Gemma 4 Unified, as the LTX-2.5 text encoder. +#' +#' @param config Model configuration list: the checkpoint's +#' \code{text_config} fields (see \code{\link{load_gemma4_text_encoder}}). +#' +#' @return Module whose forward(input_ids, attention_mask, +#' output_hidden_states) returns \code{list(last_hidden_state, +#' hidden_states)}. \code{input_ids} are 0-based; \code{hidden_states} +#' follows transformers: the state entering each layer, then the +#' final-normed output (\code{num_hidden_layers + 1} entries). +#' +#' @export +gemma4_text_model <- torch::nn_module( + "Gemma4TextModel", + initialize = function(config) { + self$config <- config + self$embed_tokens <- embedding_noinit(config$vocab_size, config$hidden_size) + self$layers <- torch::nn_module_list(lapply(seq_len(config$num_hidden_layers) - 1L, + function(i) gemma4_decoder_layer(config, i))) + self$norm <- gemma4_rms_norm(config$hidden_size, eps = config$rms_norm_eps) + rp <- config$rope_parameters + self$inv_freq <- list( + sliding_attention = .gemma4_inv_freq(config$head_dim, rp$sliding_attention), + full_attention = .gemma4_inv_freq(config$global_head_dim, rp$full_attention) + ) +}, + forward = function(input_ids, attention_mask = NULL, + output_hidden_states = FALSE) { + b <- input_ids$shape[1] + s <- input_ids$shape[2] + device <- input_ids$device + w <- self$embed_tokens$weight + # The embedding scale is rounded to the weight dtype, as in transformers + # (sqrt(3840) = 61.97 becomes 62 in bfloat16) + scale <- torch::torch_tensor(sqrt(self$config$hidden_size), dtype = w$dtype, + device = w$device) + hidden_states <- self$embed_tokens(input_ids + 1L) * scale + dtype <- hidden_states$dtype + + masks <- .gemma4_masks(attention_mask, b, s, self$config$sliding_window, + dtype, device) + positions <- torch::torch_arange(0, s - 1L, device = device)$unsqueeze(1L)$expand(c(b, -1L)) + rope <- lapply(self$inv_freq, .gemma4_rope, position_ids = positions, + dtype = dtype) + + all_hidden <- NULL + for (i in seq_along(self$layers)) { + if (output_hidden_states) { + all_hidden <- c(all_hidden, list(hidden_states)) + } + type <- self$config$layer_types[[i]] + hidden_states <- self$layers[[i]](hidden_states, masks[[type]], + rope[[type]][[1]], rope[[type]][[2]]) + } + hidden_states <- self$norm(hidden_states) + if (output_hidden_states) { + all_hidden <- c(all_hidden, list(hidden_states)) + } + list(last_hidden_state = hidden_states, hidden_states = all_hidden) +} +) + +# Additive [B, 1, S, S] masks per layer type: causal, keys limited to +# attention_mask, and for sliding layers q - k < window. Masked entries +# take the dtype's finite minimum, so a fully masked row (a left-pad +# query) softmaxes to a uniform finite row instead of NaN. +.gemma4_masks <- function(attention_mask, b, s, window, dtype, device) { + q <- torch::torch_arange(0, s - 1L, device = device)$view(c(s, 1L)) + k <- torch::torch_arange(0, s - 1L, device = device)$view(c(1L, s)) + causal <- k$le(q) + sliding <- causal$logical_and((q - k)$lt(window)) + keep <- if (is.null(attention_mask)) { + torch::torch_ones(c(b, 1L, 1L, s), dtype = torch::torch_bool(), + device = device) + } else { + attention_mask$to(device = device)$gt(0)$view(c(b, 1L, 1L, s)) + } + neg <- .gemma4_dtype_min(dtype) + additive <- function(allowed) { + allowed <- keep$logical_and(allowed$view(c(1L, 1L, s, s))) + torch::torch_zeros(c(b, 1L, s, s), dtype = dtype, device = device)$ + masked_fill(allowed$logical_not(), neg) + } + list(full_attention = additive(causal), + sliding_attention = additive(sliding)) +} + +# Most negative finite value of a float dtype. R torch's torch_finfo() +# rejects bfloat16, whose minimum is -(2 - 2^-7) * 2^127 (float32's +# minimum would round to -Inf there). +.gemma4_dtype_min <- function(dtype) { + if (identical(dtype$.type(), "BFloat16")) { + return(-(2 - 2^-7) * 2^127) + } + torch::torch_finfo(dtype)$min +} + +# Python str.strip(): trims every character str.isspace() accepts, which +# includes Unicode spaces (no-break, ideographic, ...) that base +# trimws() keeps +.py_strip <- function(x) { + ws <- intToUtf8(c(9:13, 28:32, 0x85, 0xA0, 0x1680, 0x2000:0x200A, + 0x2028, 0x2029, 0x202F, 0x205F, 0x3000), multiple = TRUE) + cls <- paste0("[", paste(ws, collapse = ""), "]") + gsub(paste0("^", cls, "+|", cls, "+$"), "", enc2utf8(x), perl = TRUE) +} + +# Model configuration from a checkpoint's text_config (the LTX-2.5 file +# embeds the full HF config as gemma_config metadata) +.gemma4_build_config <- function(raw) { + tc <- raw$text_config %||% raw + unsupported <- c(num_kv_shared_layers = 0, hidden_size_per_layer_input = 0) + for (k in names(unsupported)) { + if (!is.null(tc[[k]]) && tc[[k]] != unsupported[[k]]) { + stop("Gemma 4 option ", k, " = ", tc[[k]], " is not supported") + } + } + if (isTRUE(tc$enable_moe_block)) { + stop("Gemma 4 MoE blocks are not supported") + } + list( + vocab_size = as.integer(tc$vocab_size), + hidden_size = as.integer(tc$hidden_size), + intermediate_size = as.integer(tc$intermediate_size), + num_hidden_layers = as.integer(tc$num_hidden_layers), + num_attention_heads = as.integer(tc$num_attention_heads), + num_key_value_heads = as.integer(tc$num_key_value_heads), + num_global_key_value_heads = as.integer(tc$num_global_key_value_heads %||% + tc$num_key_value_heads), + head_dim = as.integer(tc$head_dim), + global_head_dim = as.integer(tc$global_head_dim %||% 512L), + attention_k_eq_v = isTRUE(tc$attention_k_eq_v), + layer_types = as.list(tc$layer_types), + sliding_window = as.integer(tc$sliding_window), + rope_parameters = tc$rope_parameters, + rms_norm_eps = tc$rms_norm_eps %||% 1e-6 + ) +} + +# Checkpoint key -> module name: text weights only (model.* in the LTX +# file, model.language_model.* in a transformers export); NA skips the +# vision/audio towers, projectors, the connectors' text projection and +# the embedded tokenizer assets +.gemma4_map_key <- function(key) { + if (grepl("^model\\.language_model\\.", key)) { + return(sub("^model\\.language_model\\.", "", key)) + } + if (grepl("^model\\.(embed_tokens|layers|norm)\\.", key)) { + return(sub("^model\\.", "", key)) + } + NA_character_ +} + +#' Load the Gemma 4 text encoder (LTX-2.5) +#' +#' Streams the text weights of an LTX-2.5 Gemma 4 checkpoint into a +#' \code{\link{gemma4_text_model}} skeleton, one tensor at a time, so +#' the 26 GB bf16 file is never fully in memory. The vision and audio +#' towers in the same file are skipped. +#' +#' @param path Path to the LTX-2.5 text encoder file (e.g. +#' \code{gemma4-12b-with-proj-ltx-2.5-bf16.safetensors}), which embeds +#' its config as \code{gemma_config} metadata, or an NF4 artifact +#' directory from \code{\link{gemma4_quantize_nf4}} (dispatched to +#' \code{\link{load_gemma4_nf4}}). +#' @param device Character. Device for the loaded model. +#' @param dtype Character. "bfloat16" (checkpoint native) or "float32". +#' @param pin Logical. When loading to the CPU, page-lock the weights so +#' \code{\link{encode_with_gemma4}} can move the model to the GPU per +#' encode at DMA speed. Default follows +#' \code{options(diffuseR.pin_staging)}. +#' @param verbose Logical. +#' +#' @return A \code{gemma4_text_model}. +#' +#' @export +load_gemma4_text_encoder <- function(path, device = "cpu", dtype = "bfloat16", + pin = getOption("diffuseR.pin_staging", TRUE), + verbose = TRUE) { + if (file.exists(file.path(path, "manifest.json"))) { + return(load_gemma4_nf4(path, device = device, dtype = dtype, + pin = pin, verbose = verbose)) + } + ckpt <- ltx23_open_checkpoint(path, require_version = NULL) + meta <- ckpt$handle$metadata[["__metadata__"]] + if (is.null(meta[["gemma_config"]])) { + stop("No gemma_config metadata in ", path, + "; expected an LTX-2.5 Gemma 4 text encoder file") + } + config <- .gemma4_build_config(jsonlite::fromJSON(meta[["gemma_config"]], + simplifyVector = FALSE)) + torch_dtype <- switch(dtype, bfloat16 = torch::torch_bfloat16(), + float32 = torch::torch_float32(), + stop("Unsupported dtype: ", dtype)) + if (verbose) { + message(sprintf("Creating Gemma 4 model: %d layers, hidden_size=%d", + config$num_hidden_layers, config$hidden_size)) + } + model <- .construct_skeleton(gemma4_text_model, config, dtype = torch_dtype) + + keys <- ckpt$keys + res <- ltx23_load_group(ckpt, keys, model, map_key = .gemma4_map_key, + verbose = verbose) + if (length(res$unmapped) || length(res$unfilled)) { + stop("Gemma 4 load: ", length(res$unmapped), " unmapped, ", + length(res$unfilled), " unfilled (e.g. ", + paste(utils::head(c(res$unmapped, res$unfilled), 3), collapse = ", "), + ")") + } + model$to(device = device) + model$eval() + + if (pin && device == "cpu") { + st <- .pin_component(model) + if (!is.null(st)) { + attr(model, "staging") <- st + if (verbose) { + message("Gemma 4 weights pinned for staged transfer") + } + } + } + model +} + +#' Encode prompts with Gemma 4 for LTX-2.5 +#' +#' The LTX-2.5 text-encoder step: strips each prompt (Python +#' \code{str.strip()} semantics, Unicode spaces included), tokenizes with left +#' padding to \code{max_sequence_length}, and returns every hidden state +#' (embedding output plus each layer, the last final-normed) stacked on a +#' new last dimension for the connectors. +#' +#' @param prompts Character vector. +#' @param model A \code{\link{gemma4_text_model}}. A pinned CPU-resident +#' model moves to \code{device} for the encode and back afterwards. +#' @param tokenizer A \code{\link{gemma4_tokenizer}}. +#' @param max_sequence_length Integer. Padded token length. +#' @param device Character. Compute device. +#' @param verbose Logical. +#' +#' @return List with \code{prompt_embeds} \code{[batch, seq_len, +#' hidden_size, num_layers + 1]} and \code{prompt_attention_mask} +#' \code{[batch, seq_len]}, the same layout as +#' \code{\link{encode_with_gemma3}}. +#' +#' @export +encode_with_gemma4 <- function(prompts, model, tokenizer, + max_sequence_length = 1024L, device = "cuda", + verbose = TRUE) { + stopifnot(inherits(tokenizer, "gemma4_tokenizer")) + # Same staging contract as encode_with_gemma3(): arm the offload + # before the onload so a partial onload is still undone + staging <- attr(model, "staging") + if (!is.null(staging) && device != "cpu" && !.staged_on(staging, device)) { + on.exit(.staged_offload(staging), add = TRUE) + .staged_onload(staging, device) + } + tokens <- encode_gemma4(tokenizer, .py_strip(as.character(prompts)), + max_length = max_sequence_length) + input_ids <- torch::torch_tensor(tokens$input_ids, dtype = torch::torch_long(), + device = device) + attention_mask <- torch::torch_tensor(tokens$attention_mask, + dtype = torch::torch_long(), + device = device) + if (verbose) { + message("Encoding with Gemma 4...") + } + torch::with_no_grad({ + out <- model(input_ids, attention_mask = attention_mask, + output_hidden_states = TRUE) + }) + list(prompt_embeds = torch::torch_stack(out$hidden_states, dim = -1L), + prompt_attention_mask = attention_mask) +} diff --git a/R/jit_ltx23.R b/R/jit_ltx23.R index 27c754f..d6af95c 100644 --- a/R/jit_ltx23.R +++ b/R/jit_ltx23.R @@ -44,10 +44,14 @@ def nf4_lin(x: Tensor, packed: Tensor, absmax: Tensor, bias: Tensor, table2: Ten # per half-byte weight and made dequant the per-step wall. absmax # is cast down before the scale so the product never promotes to # float32 (which would materialize a full-precision weight copy). - rows = bias.size(0) + # A zero-length bias is the no-bias sentinel (List[Tensor] slots + # cannot hold None); the LTX-2.5 video FF has none. vals = torch.embedding(table2, packed.int()) w = vals.reshape([-1, 64]) * absmax.type_as(table2).unsqueeze(1) - return torch.linear(x, w.reshape([rows, -1]).type_as(x), bias) + w = w.reshape([-1, x.size(-1)]).type_as(x) + if bias.numel() == 0: + return torch.linear(x, w) + return torch.linear(x, w, bias) def rmsn(x: Tensor) -> Tensor: v = x.float().pow(2).mean(-1, keepdim=True) @@ -218,9 +222,16 @@ def stack_nf4(h: Tensor, ah: Tensor, enc: Tensor, aenc: Tensor, } .ltx23_jit_pack_ff <- function(ff) { - list(ff$net[[1]]$proj$weight_nf4, ff$net[[1]]$proj$weight_absmax, - ff$net[[1]]$proj$bias, ff$net[[3]]$weight_nf4, - ff$net[[3]]$weight_absmax, ff$net[[3]]$bias) + proj <- ff$net[[1]]$proj + out <- ff$net[[3]] + list(proj$weight_nf4, proj$weight_absmax, .ltx23_jit_bias(proj), + out$weight_nf4, out$weight_absmax, .ltx23_jit_bias(out)) +} + +# A linear's bias, or the zero-length no-bias sentinel nf4_lin expects +.ltx23_jit_bias <- function(lin) { + lin$bias %||% torch::torch_empty(0L, dtype = lin$weight_absmax$dtype, + device = lin$weight_absmax$device) } #' Pack a transformer block's weights for the JIT stack diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R index ff47354..0ba3a36 100644 --- a/R/nf4_ltx23.R +++ b/R/nf4_ltx23.R @@ -255,7 +255,12 @@ ltx23_nf4_linear <- torch::nn_module( #' Non-cast tensors are copied through unchanged. The manifest carries #' \code{format = "nf4"}. #' -#' @param checkpoint_path Source .safetensors (bf16 single file). +#' @param checkpoint_path Source .safetensors (bf16 single file), or an +#' open \code{ltx23_checkpoint}. For LTX-2.5, pass +#' \code{\link{ltx25_open_checkpoint}} to write one self-contained +#' artifact (transformer, connectors and text projection, VAEs, +#' vocoder); when its transformer is already an NF4 artifact, the +#' NF4 tensors are bundled as they are, with no re-quantization. #' @param output_dir Output directory for shards + manifest; NULL (the #' default) resolves under tools::R_user_dir("diffuseR", "data"). #' @param shard_bytes Numeric. Target shard size in bytes. The default @@ -289,7 +294,19 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL, } dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) - ckpt <- ltx23_open_checkpoint(checkpoint_path) + if (inherits(checkpoint_path, "ltx23_checkpoint")) { + # An open checkpoint, e.g. the LTX-2.5 file set from + # ltx25_open_checkpoint(); one already in NF4 is bundled as is + ckpt <- checkpoint_path + source <- basename(unname(ckpt$sources %||% ckpt$path)) + source <- source[!is.na(source)] + } else { + ckpt <- ltx23_open_checkpoint(checkpoint_path) + source <- basename(checkpoint_path) + } + passthrough <- identical(ckpt$format, "nf4") + prefix <- sprintf("ltx%s-nf4", + sub("^(\\d+\\.\\d+).*", "\\1", ckpt$version %||% "2.3")) shard <- list() shard_size <- 0 @@ -300,7 +317,7 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL, if (!length(shard)) { return() } - fname <- sprintf("ltx2.3-nf4-%05d.safetensors", + fname <- sprintf("%s-%05d.safetensors", prefix, length(shard_files) + 1L) safetensors::safe_save_file(shard, file.path(output_dir, fname)) shard_files[[length(shard_files) + 1L]] <<- fname @@ -319,7 +336,11 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL, tensor <- ckpt$handle$get_tensor(key) mapped <- ltx23_map_dit_key(key) - if (startsWith(key, "model.diffusion_model.") && + if (passthrough) { + shard[[key]] <- tensor + shard_size <- shard_size + prod(tensor$shape) * tensor$element_size() + n_cast <- n_cast + endsWith(key, "_absmax") + } else if (startsWith(key, "model.diffusion_model.") && ltx23_is_fp8_cast_key(mapped)) { torch::with_no_grad({ q <- ltx23_nf4_quantize(tensor) @@ -337,17 +358,19 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL, if (shard_size >= shard_bytes) { flush_shard() } - if (i %% 200L == 0L) { - gc(verbose = FALSE) - if (verbose) { - message(sprintf(" quantizing %d/%d tensors", i, length(keys))) - } + # Per tensor: R's gc cannot see torch's native memory, and 200 + # dropped bf16 sources plus float32 quantization temporaries + # (~0.27 GB each for the FFN weights) outgrow the free RAM of a + # host sharing memory with other services (cf. gemma4_quantize_nf4) + gc(verbose = FALSE) + if (i %% 200L == 0L && verbose) { + message(sprintf(" quantizing %d/%d tensors", i, length(keys))) } } flush_shard() manifest <- list( - source = basename(checkpoint_path), + source = source, model_version = ckpt$version, format = "nf4", shards = shard_files, @@ -382,7 +405,8 @@ ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = NULL, ltx23_load_transformer_nf4 <- function(ckpt, device = "cuda", verbose = TRUE, ...) { stopifnot(inherits(ckpt, "ltx23_checkpoint")) - model <- .construct_skeleton(ltx23_transformer, ...) + args <- utils::modifyList(ltx23_dit_config_args(ckpt$config), list(...)) + model <- do.call(.construct_skeleton, c(list(ltx23_transformer), args)) groups <- ltx23_split_keys(ckpt$keys) dit_keys <- groups$dit diff --git a/R/quantize_gemma4.R b/R/quantize_gemma4.R new file mode 100644 index 0000000..af604c6 --- /dev/null +++ b/R/quantize_gemma4.R @@ -0,0 +1,258 @@ +#' Quantize the LTX-2.5 Gemma 4 text encoder to NF4 shards +#' +#' Streams the LTX-2.5 encoder file tensor by tensor. The per-layer +#' projection weights (q/k/v/o and gate/up/down; 328 tensors in the 12B, +#' since full-attention layers have no v_proj) are stored as NF4 (packed +#' uint8 + \code{_absmax} float32 blocks); embeddings, norms and the +#' per-layer \code{layer_scalar} are copied at the resident dtype. Vision +#' and audio towers are dropped. +#' +#' The artifact is self-contained for the LTX-2.5 text side: it also +#' keeps the connectors' \code{text_embedding_projection.*} weights +#' (resident dtype, original keys) and writes the embedded tokenizer as +#' \code{tokenizer.json}, so the 26 GB source file is not needed after +#' quantizing. Load it with \code{\link{load_gemma4_nf4}} (or +#' \code{\link{load_gemma4_text_encoder}} on the directory) and the +#' tokenizer with \code{gemma4_tokenizer(output_dir)}. +#' +#' Peak host memory is about one shard (\code{shard_bytes}) plus the +#' largest single source tensor and its save buffer (the 2 GB bf16 +#' embedding); measured under a 20 GB cgroup cap on the 12B. +#' +#' @param path The LTX-2.5 text encoder file +#' (\code{gemma4-12b-with-proj-ltx-2.5-bf16.safetensors}). +#' @param output_dir Output directory (default: \code{gemma4-nf4} under +#' \code{tools::R_user_dir}). +#' @param shard_bytes Numeric. Target shard size in bytes; the 1.9e9 +#' default keeps shards readable by safetensors older than 0.3.0. +#' @param force Logical. Re-quantize even if a valid manifest exists. +#' @param verbose Logical. +#' +#' @return Invisibly, the manifest list. +#' +#' @export +gemma4_quantize_nf4 <- function(path, output_dir = NULL, shard_bytes = 1.9e9, + force = FALSE, verbose = TRUE) { + if (is.null(output_dir)) { + output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + "gemma4-nf4") + } + manifest_path <- file.path(output_dir, "manifest.json") + if (!force && file.exists(manifest_path)) { + manifest <- jsonlite::fromJSON(manifest_path) + # Complete means every shard plus the tokenizer the manifest + # promises; anything less is rebuilt + if (identical(manifest$model, "gemma4") && + identical(manifest$format, "nf4") && + all(file.exists(file.path(output_dir, + c(manifest$shards, manifest$tokenizer))))) { + if (verbose) { + message("NF4 artifact already present: ", output_dir) + } + return(invisible(manifest)) + } + } + + ckpt <- ltx23_open_checkpoint(path, require_version = NULL) + meta <- ckpt$handle$metadata[["__metadata__"]] + if (is.null(meta[["gemma_config"]])) { + stop("No gemma_config metadata in ", path, + "; expected an LTX-2.5 Gemma 4 text encoder file") + } + config_raw <- jsonlite::fromJSON(meta[["gemma_config"]], + simplifyVector = FALSE) + resident_dtype <- if (.st_can_write("bfloat16")) { + torch::torch_bfloat16() + } else { + message("safetensors cannot write bfloat16; storing resident ", + "tensors as float32 (larger artifact, same results)") + torch::torch_float32() + } + resident_bytes <- torch::torch_empty(1L, dtype = resident_dtype)$element_size() + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + # A stale manifest must not validate a half-written artifact + unlink(manifest_path) + + if ("tokenizer_json" %in% ckpt$keys) { + writeLines(.gemma4_tokenizer_json(path), + file.path(output_dir, "tokenizer.json"), useBytes = TRUE) + } + + shard <- list() + shard_size <- 0 + shard_files <- character(0) + n_cast <- 0L + flush_shard <- function() { + if (!length(shard)) { + return() + } + fname <- sprintf("gemma4-nf4-%05d.safetensors", length(shard_files) + 1L) + safetensors::safe_save_file(shard, file.path(output_dir, fname)) + shard_files[[length(shard_files) + 1L]] <<- fname + if (verbose) { + message(sprintf(" wrote %s (%.2f GB, %d tensors)", fname, + shard_size / 1e9, length(shard))) + } + shard <<- list() + shard_size <<- 0 + gc(verbose = FALSE) + } + + for (key in ckpt$keys) { + name <- .gemma4_map_key(key) + if (is.na(name)) { + if (!startsWith(key, "text_embedding_projection.")) { + next + } + name <- key # kept for the LTX-2.5 connectors + } + tensor <- ckpt$handle$get_tensor(key) + quant <- .gemma3_is_quant_key(name) + bytes <- prod(tensor$shape) * if (quant) 0.5625 else resident_bytes + # Flush first so a shard stays under shard_bytes; a tensor larger + # than the target alone (the 2 GB embedding) gets its own shard + if (shard_size + bytes > shard_bytes) { + flush_shard() + } + if (quant) { + torch::with_no_grad({ + q <- ltx23_nf4_quantize(tensor) + shard[[name]] <- q$packed + shard[[paste0(name, "_absmax")]] <- q$absmax + }) + n_cast <- n_cast + 1L + } else { + shard[[name]] <- tensor$to(dtype = resident_dtype) + } + shard_size <- shard_size + bytes + rm(tensor) + # R's gc cannot see torch's native memory: without a collection + # per tensor the dropped source tensors and quantization + # temporaries (up to ~0.24 GB each) pile up past 20 GB within + # seconds on the 12B (measured: cgroup OOM kill) + gc(verbose = FALSE) + } + flush_shard() + + manifest <- list(model = "gemma4", format = "nf4", shards = shard_files, + n_cast = n_cast, config = config_raw) + if (file.exists(file.path(output_dir, "tokenizer.json"))) { + manifest$tokenizer <- "tokenizer.json" + } + jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE, + pretty = TRUE, digits = NA, null = "null") + if (verbose) { + message(sprintf("Quantized %d weights to NF4 across %d shards: %s", + n_cast, length(shard_files), output_dir)) + } + invisible(manifest) +} + +#' Load the LTX-2.5 Gemma 4 text encoder from an NF4 artifact +#' +#' Builds the model as a skeleton at the compute dtype, swaps the +#' projection linears for NF4 modules filled from the artifact +#' (dequantized per forward), copies the residents (parameters and the +#' \code{layer_scalar} buffers), and hard-errors on anything the artifact +#' does not fill. The connectors' projection weights in the artifact are +#' left for the connector loader. +#' +#' @param artifact_dir Directory produced by \code{\link{gemma4_quantize_nf4}}. +#' @param device "cuda" or "cpu". +#' @param dtype Compute dtype ("bfloat16" default, or "float32"). +#' @param pin Logical. When loading to the CPU, page-lock the weights so +#' \code{\link{encode_with_gemma4}} can move the model to the GPU per +#' encode at DMA speed. Default follows +#' \code{options(diffuseR.pin_staging)}. +#' @param verbose Logical. +#' +#' @return A \code{gemma4_text_model}. +#' +#' @export +load_gemma4_nf4 <- function(artifact_dir, device = "cuda", dtype = "bfloat16", + pin = getOption("diffuseR.pin_staging", TRUE), + verbose = TRUE) { + ckpt <- ltx23_open_fp8_checkpoint(artifact_dir) + manifest <- jsonlite::fromJSON(file.path(artifact_dir, "manifest.json"), + simplifyVector = FALSE) + if (!identical(manifest$model, "gemma4")) { + stop("Not a Gemma 4 NF4 artifact (manifest model is ", + manifest$model %||% "absent", "): ", artifact_dir) + } + config <- .gemma4_build_config(manifest$config) + torch_dtype <- switch(dtype, bfloat16 = torch::torch_bfloat16(), + float32 = torch::torch_float32(), + stop("Unsupported dtype: ", dtype)) + if (verbose) { + message(sprintf("Creating Gemma 4 NF4 model: %d layers, hidden_size=%d", + config$num_hidden_layers, config$hidden_size)) + } + model <- .construct_skeleton(gemma4_text_model, config, dtype = torch_dtype) + + cast_keys <- sub("_absmax$", "", grep("_absmax$", ckpt$keys, value = TRUE)) + torch::with_no_grad({ + for (key in cast_keys) { + # layers....weight + segs <- strsplit(sub("\\.weight$", "", key), ".", fixed = TRUE)[[1]] + parent <- model$layers[[as.integer(segs[2]) + 1L]][[segs[3]]] + old <- parent[[segs[4]]] + nf4 <- ltx23_nf4_linear(old$weight$shape[1], old$weight$shape[2], + bias = FALSE) + nf4$set_nf4_weight(ckpt$handle$get_tensor(key), + ckpt$handle$get_tensor(paste0(key, "_absmax"))) + do.call(`$<-`, list(parent, segs[4], nf4)) + } + + dests <- c(model$named_parameters(), model$named_buffers()) + dests <- dests[!grepl("weight_(nf4|absmax)$", names(dests))] + filled <- character(0) + residents <- setdiff(ckpt$keys, c(cast_keys, paste0(cast_keys, "_absmax"))) + for (key in residents) { + dest <- dests[[key]] + if (is.null(dest)) { + next # e.g. text_embedding_projection.*, for the connectors + } + dest$copy_(ckpt$handle$get_tensor(key)) + filled <- c(filled, key) + } + unfilled <- setdiff(names(dests), filled) + if (length(unfilled)) { + stop("Gemma 4 NF4 load: ", length(unfilled), + " parameters not filled, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + }) + + model$to(device = device) + model$eval() + if (pin && device == "cpu") { + st <- .pin_component(model) + if (!is.null(st)) { + attr(model, "staging") <- st + if (verbose) { + message("Gemma 4 weights pinned for staged transfer") + } + } + } + model +} + +# TRUE when `path` is a Gemma 4 encoder: its NF4 artifact directory, or the +# LTX-2.5 single file (gemma_config metadata). Lets a caller holding a path +# pick the Gemma 4 loader over the Gemma 3 one. +.is_gemma4_artifact <- function(path) { + if (!is.character(path) || length(path) != 1L) { + return(FALSE) + } + manifest <- file.path(path, "manifest.json") + if (file.exists(manifest)) { + return(identical(jsonlite::fromJSON(manifest)$model, "gemma4")) + } + if (!file.exists(path) || dir.exists(path)) { + return(FALSE) + } + meta <- tryCatch( + ltx23_open_checkpoint(path, require_version = NULL)$handle$metadata[["__metadata__"]], + error = function(e) NULL) + !is.null(meta[["gemma_config"]]) +} diff --git a/R/resident.R b/R/resident.R index 0ace6d7..d46f6be 100644 --- a/R/resident.R +++ b/R/resident.R @@ -236,7 +236,9 @@ #' \code{\link{download_sdxl}} cache). #' @param verbose Print progress messages. #' @param text_encoder,tokenizer \code{ltx} only: paths to the Gemma3 encoder -#' artifact and the tokenizer directory. \code{\link{ltx23_load_pipeline}} +#' artifact and the tokenizer directory, or for LTX-2.5 the Gemma 4 NF4 +#' artifact (\code{\link{gemma4_quantize_nf4}}) for both, since it carries +#' its tokenizer. \code{\link{ltx23_load_pipeline}} #' does not load these -- \code{\link{txt2vid_ltx2}} takes them per call -- #' so a handle built without them can be activated and cannot generate. #' Given here they are loaded ONCE, pinned on the host, and passed to every @@ -346,10 +348,19 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", stop("text_encoder needs a tokenizer: the encode takes both", call. = FALSE) } - if (verbose) message("Loading the Gemma3 text encoder (pinned)...") - res$text_encoder <- load_gemma3_text_encoder( - text_encoder, device = "cpu", pin = TRUE, verbose = verbose) - res$tokenizer <- gemma3_tokenizer(tokenizer) + # LTX-2.5 encodes with Gemma 4; its NF4 artifact carries the + # tokenizer too. txt2vid_ltx2() picks the encode by tokenizer class. + if (.is_gemma4_artifact(text_encoder)) { + if (verbose) message("Loading the Gemma 4 text encoder (pinned)...") + res$text_encoder <- load_gemma4_text_encoder( + text_encoder, device = "cpu", pin = TRUE, verbose = verbose) + res$tokenizer <- gemma4_tokenizer(tokenizer) + } else { + if (verbose) message("Loading the Gemma3 text encoder (pinned)...") + res$text_encoder <- load_gemma3_text_encoder( + text_encoder, device = "cpu", pin = TRUE, verbose = verbose) + res$tokenizer <- gemma3_tokenizer(tokenizer) + } ## Counted, so `resident_status()` reports what the process actually ## holds. A pinned set omitted from the total reads as headroom that ## is not there, and the fleet's admission arithmetic is downstream diff --git a/R/tokenizer_gemma4.R b/R/tokenizer_gemma4.R new file mode 100644 index 0000000..8ab8701 --- /dev/null +++ b/R/tokenizer_gemma4.R @@ -0,0 +1,221 @@ +#' Gemma 4 SentencePiece-style BPE tokenizer +#' +#' Native R port of the HuggingFace \code{tokenizers} pipeline in the +#' Gemma 4 \code{tokenizer.json} shipped with LTX-2.5: added-token +#' literals split out first, the \code{Replace " " -> U+2581} normalizer, +#' characters mapped to vocab ids with \code{<0xHH>} byte fallback, then +#' rank-ordered BPE merges over the whole string (the \code{Split} on +#' " " pre-tokenizer never matches after the normalizer). The +#' post-processor template decides special tokens; for Gemma 4 it adds +#' none, so there is no BOS. +#' +#' @name tokenizer_gemma4 +NULL + +#' Load the Gemma 4 tokenizer +#' +#' @param source Path to a \code{tokenizer.json}, a directory holding +#' one, or an LTX-2.5 text encoder \code{.safetensors} file, which +#' carries the same JSON in its \code{tokenizer_json} tensor. +#' +#' @return A \code{gemma4_tokenizer}. +#' +#' @export +gemma4_tokenizer <- function(source) { + tj <- jsonlite::fromJSON(.gemma4_tokenizer_json(source), + simplifyVector = TRUE) + model <- tj$model + if (!identical(model$type, "BPE")) { + stop("Expected a BPE tokenizer, got: ", model$type %||% "none") + } + norm <- tj$normalizer + if (!identical(norm$type, "Replace") || is.null(norm$pattern$String)) { + stop("Unsupported normalizer: ", norm$type %||% "none", + " (expected a Replace of a literal string)") + } + + # Same integer-id layout as qwen_bpe_tokenizer(): atomic vectors only, + # since live environment bindings make every R gc slower + vocab <- unlist(model$vocab) # named int: piece -> id (0-based) + piece_names <- names(vocab) + merges <- model$merges + if (is.matrix(merges)) { + a <- merges[, 1] + b <- merges[, 2] + } else { + sp <- regexpr(" ", merges, fixed = TRUE) + a <- substr(merges, 1L, sp - 1L) + b <- substr(merges, sp + 1L, nchar(merges)) + } + id_a <- unname(vocab[match(a, piece_names)]) + id_b <- unname(vocab[match(b, piece_names)]) + id_r <- unname(vocab[match(paste0(a, b), piece_names)]) + ok <- !is.na(id_a) & !is.na(id_b) & !is.na(id_r) + key <- id_a[ok] * 2097152 + id_b[ok] # ids < 2^21: exact in doubles + ord <- order(key) + + byte_pieces <- sprintf("<0x%02X>", 0:255) + added <- tj$added_tokens + + structure( + list( + merge_key = key[ord], + merge_rank = seq_along(id_a)[ok][ord], + merge_result = id_r[ok][ord], + piece_names = piece_names, + piece_ids = unname(vocab), + byte_ids = unname(vocab[match(byte_pieces, piece_names)]), + byte_fallback = isTRUE(model$byte_fallback), + unk_id = unname(vocab[model$unk_token %||% ""]), + norm_from = norm$pattern$String, + norm_to = norm$content, + added = list2env(stats::setNames(as.list(added$id), + added$content), + parent = emptyenv()), + added_contents = added$content[order(-nchar(added$content))], + template = .gemma4_template_ids(tj$post_processor, added), + pad_id = added$id[match("", added$content)] %||% 0L + ), + class = "gemma4_tokenizer" + ) +} + +#' @export +print.gemma4_tokenizer <- function(x, ...) { + cat("\n") + cat(" vocab: ", length(x$piece_ids), "pieces,", + length(x$merge_key), "merges\n") + cat(" special:", length(x$template$prefix), "prefix /", + length(x$template$suffix), "suffix ids per sequence\n") + invisible(x) +} + +# The tokenizer.json text from a json path, a directory, or the +# tokenizer_json uint8 tensor of an LTX-2.5 encoder file +.gemma4_tokenizer_json <- function(source) { + path <- path.expand(source) + if (dir.exists(path)) { + path <- file.path(path, "tokenizer.json") + } + if (!file.exists(path)) { + stop("Tokenizer source not found: ", source) + } + if (!grepl("\\.safetensors$", path)) { + return(paste(readLines(path, warn = FALSE, encoding = "UTF-8"), + collapse = "\n")) + } + handle <- safetensors::safetensors$new(path, framework = "torch") + if (!"tokenizer_json" %in% handle$keys()) { + stop("No tokenizer_json tensor in ", path) + } + bytes <- as.raw(as.integer(as.array(handle$get_tensor("tokenizer_json")))) + txt <- rawToChar(bytes) + Encoding(txt) <- "UTF-8" + txt +} + +# Special-token ids the TemplateProcessing "single" template puts before +# and after the sequence +.gemma4_template_ids <- function(post, added) { + out <- list(prefix = integer(0), suffix = integer(0)) + single <- post$single + if (!identical(post$type, "TemplateProcessing") || is.null(single)) { + return(out) + } + seen_seq <- FALSE + for (i in seq_len(NROW(single))) { + sp <- single$SpecialToken$id[i] %||% NA + if (!is.null(single$Sequence) && !is.na(single$Sequence$id[i] %||% NA)) { + seen_seq <- TRUE + } else if (!is.na(sp)) { + id <- added$id[match(sp, added$content)] + if (seen_seq) { + out$suffix <- c(out$suffix, id) + } else { + out$prefix <- c(out$prefix, id) + } + } + } + out +} + +# Normalize one plain-text segment, map characters to ids (byte fallback +# for pieces missing from the vocab), then apply the ranked merges. +# Shares the merge loop with the Qwen tokenizer. +.gemma4_encode_segment <- function(text, tokenizer) { + if (!nzchar(text)) { + return(integer(0)) + } + text <- gsub(tokenizer$norm_from, tokenizer$norm_to, enc2utf8(text), + fixed = TRUE) + chars <- strsplit(text, "", fixed = TRUE)[[1]] + ids <- tokenizer$piece_ids[match(chars, tokenizer$piece_names)] + if (anyNA(ids)) { + parts <- as.list(ids) + for (i in which(is.na(ids))) { + if (tokenizer$byte_fallback) { + b <- as.integer(charToRaw(chars[i])) + parts[[i]] <- tokenizer$byte_ids[b + 1L] + } else { + parts[[i]] <- tokenizer$unk_id + } + } + ids <- unlist(parts) + } + as.integer(.qwen_bpe_merge_ids(ids, tokenizer)) +} + +#' Encode prompts with the Gemma 4 tokenizer +#' +#' Tokenizes, applies the post-processor's special tokens (none for +#' Gemma 4), truncates to \code{max_length} keeping the start, and pads +#' on the left with \code{}, as the LTX pipelines do. +#' +#' @param tokenizer A \code{\link{gemma4_tokenizer}}. +#' @param texts Character vector of prompts. +#' @param max_length Integer. Fixed sequence length, or NULL for ragged +#' unpadded output. +#' +#' @return List with \code{input_ids} and \code{attention_mask}: integer +#' matrices [length(texts), max_length] (ragged lists when +#' \code{max_length} is NULL). Ids are 0-based. +#' +#' @export +encode_gemma4 <- function(tokenizer, texts, max_length = 1024L) { + stopifnot(inherits(tokenizer, "gemma4_tokenizer")) + encode_one <- function(text) { + ids <- integer(0) + for (chunk in .qwen_split_added(text, tokenizer)) { + if (!is.na(chunk$id)) { + ids <- c(ids, as.integer(chunk$id)) + } else { + ids <- c(ids, .gemma4_encode_segment(chunk$text, tokenizer)) + } + } + c(tokenizer$template$prefix, ids, tokenizer$template$suffix) + } + all_ids <- lapply(as.character(texts), encode_one) + + if (is.null(max_length)) { + return(list( + input_ids = all_ids, + attention_mask = lapply(all_ids, function(x) rep(1L, length(x))) + )) + } + n <- length(all_ids) + input_ids <- matrix(as.integer(tokenizer$pad_id), nrow = n, ncol = max_length) + attention_mask <- matrix(0L, nrow = n, ncol = max_length) + for (i in seq_len(n)) { + ids <- all_ids[[i]] + if (length(ids) > max_length) { + ids <- ids[seq_len(max_length)] + } + len <- length(ids) + if (len > 0L) { + cols <- (max_length - len + 1L):max_length + input_ids[i, cols] <- ids + attention_mask[i, cols] <- 1L + } + } + list(input_ids = input_ids, attention_mask = attention_mask) +} diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 7767cc4..ca3d164 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -77,18 +77,66 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { latents * std + mean } -# Joint audio/video Euler denoise loop over a sigma schedule (CFG-free). +# Sampler eta for a pipeline: LTX-2.5 and later are distilled for the +# ancestral sampler (eta 1), LTX-2.3 for plain Euler (eta 0). "auto" +# with no recorded version keeps the 2.3 behavior. +.ltx23_sampler_eta <- function(sampler, model_version) { + if (sampler == "auto") { + v <- numeric_version(model_version %||% "0", strict = FALSE) + ancestral <- !is.na(v) && v >= "2.5" + sampler <- if (ancestral) "euler_ancestral" else "euler" + } + if (sampler == "euler_ancestral") 1 else 0 +} + +# One rectified-flow sampler step from sigma to sigma_next, in float32 +# (diffusers LTXEulerAncestralRFScheduler.step, Apache-2.0). eta = 0 is +# the plain Euler step x + (sigma_next - sigma) * v. eta > 0 steps +# deterministically to sigma_down <= sigma_next, then renoises back up +# to sigma_next with the signal rescaled by alpha_next / alpha_down +# (alpha = 1 - sigma), which keeps the transition variance-preserving. +# The last step (sigma_next = 0) returns the denoised estimate. +.ltx23_step <- function(x, velocity, sigma, sigma_next, eta = 0, + noise = NULL, s_noise = 1) { + if (eta == 0) { + return(x + velocity * (sigma_next - sigma)) + } + denoised <- x - velocity * sigma + if (abs(sigma_next) < 1e-8) { + return(denoised) + } + sigma_down <- sigma_next * (1 + (sigma_next / sigma - 1) * eta) + alpha_next <- 1 - sigma_next + alpha_down <- 1 - sigma_down + ratio <- sigma_down / sigma + x <- x * ratio + denoised * (1 - ratio) + renoise <- sqrt(max(sigma_next^2 - + sigma_down^2 * alpha_next^2 / (alpha_down^2 + 1e-12), 0)) + x * (alpha_next / (alpha_down + 1e-12)) + noise * (renoise * s_noise) +} + +# Joint audio/video denoise loop over a sigma schedule (CFG-free). # With a conditioning_mask [B, S] (prefix conditioning), conditioned # video tokens see a per-token timestep of zero and are frozen through -# the Euler updates (reference i2v semantics at strength 1). +# the updates (reference i2v semantics at strength 1). eta > 0 selects +# the ancestral step (LTX-2.5); its per-step noise (video, then audio) +# comes from noise_seed when given, with the caller's RNG state restored +# afterwards so later draws do not depend on the step count. .ltx23_denoise <- function(transformer, latents, audio_latents, sigmas, video_text_embeds, audio_text_embeds, text_mask, latent_frames, latent_height, latent_width, audio_num_frames, frame_rate, device, compute_dtype, verbose = TRUE, stage = "", conditioning_mask = NULL, - audio_conditioned = FALSE) { + audio_conditioned = FALSE, eta = 0, + noise_seed = NULL) { f32 <- torch::torch_float32() + if (eta > 0 && !is.null(noise_seed)) { + torch::local_torch_manual_seed(noise_seed) + } + step_noise <- function(like) { + if (eta > 0) torch::torch_randn_like(like) else NULL + } keep_mask <- if (!is.null(conditioning_mask)) { conditioning_mask$unsqueeze(3L) # [B, S, 1] for latent blending } else { @@ -161,9 +209,10 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { cond_token_index = cond_idx ) - # Euler velocity step in float32; dt is negative (sigma decreasing) - dt <- torch::torch_tensor(sigmas[i + 1L] - sigma, device = device, dtype = f32) - stepped <- latents + dt * out$sample$to(dtype = f32) + # Sampler step in float32 (latents are float32 throughout) + sigma_next <- sigmas[i + 1L] + stepped <- .ltx23_step(latents, out$sample$to(dtype = f32), sigma, + sigma_next, eta, step_noise(latents)) latents <- if (is.null(keep_mask)) { stepped } else { @@ -171,7 +220,10 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { latents * keep_mask + stepped * (1 - keep_mask) } if (!audio_conditioned) { - audio_latents <- audio_latents + dt * out$audio_sample$to(dtype = f32) + audio_latents <- .ltx23_step(audio_latents, + out$audio_sample$to(dtype = f32), + sigma, sigma_next, eta, + step_noise(audio_latents)) } rm(out) gc(verbose = FALSE) @@ -195,7 +247,9 @@ ltx23_unpack_audio_latents <- function(latents, num_mel_bins) { #' \code{ltx-2.3-22b-distilled-1.1.safetensors}) or to an fp8 artifact #' directory produced by \code{\link{ltx23_quantize_fp8}}. With the fp8 #' artifact, the transformer loads with CPU-resident fp8 weights that -#' stream to \code{device} during the forward pass. +#' stream to \code{device} during the forward pass. An open +#' \code{ltx23_checkpoint} is also accepted, such as the LTX-2.5 +#' component set from \code{\link{ltx25_open_checkpoint}}. #' @param device Character. Device for the small components (VAEs, #' vocoder, connectors) and, with fp8, the transformer residents. #' @param dtype Character. "bfloat16" (checkpoint native) or "float32". @@ -225,11 +279,17 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", } else { component_device <- device } - fp8 <- dir.exists(checkpoint_path) - ckpt <- if (fp8) { - ltx23_open_fp8_checkpoint(checkpoint_path) + if (inherits(checkpoint_path, "ltx23_checkpoint")) { + # An already-open checkpoint, e.g. ltx25_open_checkpoint() + ckpt <- checkpoint_path + fp8 <- !is.null(ckpt$format) } else { - ltx23_open_checkpoint(checkpoint_path) + fp8 <- dir.exists(checkpoint_path) + ckpt <- if (fp8) { + ltx23_open_fp8_checkpoint(checkpoint_path) + } else { + ltx23_open_checkpoint(checkpoint_path) + } } if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { # Must be set before the first CUDA allocation. NF4 runs the @@ -278,7 +338,8 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", float32 = torch::torch_float32(), stop("Unsupported dtype: ", dtype)) - pipe <- list(config = ckpt$config, checkpoint_path = ckpt$path) + pipe <- list(config = ckpt$config, checkpoint_path = ckpt$path, + model_version = ckpt$version) load_component <- function(name, module, map_key, dev) { if (verbose) { @@ -306,9 +367,9 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", ckpt, device = component_device, pin = pin, verbose = verbose ) } else { - pipe$transformer <- load_component( - "dit", .construct_skeleton(ltx23_transformer, dtype = torch_dtype), - ltx23_map_dit_key, transformer_device + pipe$transformer <- .ltx23_load_transformer( + ckpt, dtype = torch_dtype, device = transformer_device, + verbose = verbose ) } if (!is.null(attn_chunk)) { @@ -396,6 +457,12 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", if (!is.null(attr(text_encoder, "staging")) && grepl("^cuda", device)) { return(device) } + # An encoder already loaded on the GPU encodes there + on_gpu <- tryCatch(text_encoder$parameters[[1]]$device$type == "cuda", + error = function(e) FALSE) + if (isTRUE(on_gpu) && grepl("^cuda", device)) { + return(device) + } "cpu" } @@ -412,8 +479,9 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", #' \code{\link{ltx23_load_pipeline}}. #' @param text_encoder,tokenizer Gemma3 model and tokenizer (or paths; #' see \code{\link{load_gemma3_text_encoder}} and -#' \code{\link{gemma3_tokenizer}}). Ignored when \code{prompt_embeds} -#' is supplied. +#' \code{\link{gemma3_tokenizer}}), or for LTX-2.5 the Gemma 4 model +#' and \code{\link{gemma4_tokenizer}} (selected by the tokenizer's +#' class). Ignored when \code{prompt_embeds} is supplied. #' @param prompt_embeds Optional precomputed list with #' \code{prompt_embeds} (raw stacked Gemma3 states) and #' \code{prompt_attention_mask}; bypasses the text encoder. @@ -433,6 +501,12 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", #' @param guidance_scale Numeric. Only 1 (no CFG) is supported; the #' distilled checkpoints are trained for CFG-free sampling. #' @param seed Integer or NULL. +#' @param sampler One of "auto", "euler" or "euler_ancestral". "auto" +#' picks the sampler the checkpoint was distilled for: plain Euler +#' for LTX-2.3, ancestral Euler (eta 1) for LTX-2.5 and later. The +#' ancestral steps draw their noise from \code{seed + 10000} (stage 1) +#' and \code{seed + 20000} (stage 2), so the initial noise matches the +#' Euler run at the same seed. #' @param device Character. Compute device for the denoising loop. #' @param dtype Character. Model compute dtype ("bfloat16" or "float32"). #' @param filename Character or NULL. Output video path (.mp4). Audio is @@ -517,15 +591,19 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, image = NULL, condition_video = NULL, conditioning_frames = 9L, cond_noise_scale = 0, condition_latents = NULL, resident = character(), - trim_frames = 0L, audio = NULL, verbose = TRUE) { + trim_frames = 0L, audio = NULL, verbose = TRUE, + sampler = c("auto", "euler", "euler_ancestral")) { level <- .verbosity(verbose) verbose <- level == "steps" if (level != "silent") { - message(sprintf("LTX-2.3: %d frames @ %dx%d, %d steps%s", num_frames, + message(sprintf("LTX-%s: %d frames @ %dx%d, %d steps%s", + sub("^(\\d+\\.\\d+).*", "\\1", pipeline$model_version %||% "2.3"), + num_frames, width, height, length(sigmas) - 1L, if (two_stage) " (two-stage)" else "")) } stopifnot(inherits(pipeline, "ltx23_pipeline")) + eta <- .ltx23_sampler_eta(match.arg(sampler), pipeline$model_version) if (guidance_scale != 1) { stop("Only guidance_scale = 1 is supported (distilled checkpoints).") } @@ -583,7 +661,13 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, if (verbose) { message("Encoding prompt...") } - prompt_embeds <- encode_with_gemma3( + # LTX-2.5 encodes with Gemma 4 (chosen by its tokenizer class) + encode <- if (inherits(tokenizer, "gemma4_tokenizer")) { + encode_with_gemma4 + } else { + encode_with_gemma3 + } + prompt_embeds <- encode( prompt, model = text_encoder, tokenizer = tokenizer, max_sequence_length = max_sequence_length, @@ -812,7 +896,9 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, device, compute_dtype, verbose = level, stage = if (two_stage) "stage 1 " else "", conditioning_mask = conditioning_mask, - audio_conditioned = audio_conditioned + audio_conditioned = audio_conditioned, + eta = eta, + noise_seed = if (!is.null(seed)) seed + 10000 ) latents <- denoised$latents audio_latents <- denoised$audio_latents @@ -868,7 +954,9 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, video_text_embeds, audio_text_embeds, text_mask, latent_frames, latent_height, latent_width, audio_num_frames, frame_rate, - device, compute_dtype, verbose = level, stage = "stage 2 " + device, compute_dtype, verbose = level, stage = "stage 2 ", + eta = eta, + noise_seed = if (!is.null(seed)) seed + 20000 ) latents <- denoised$latents audio_latents <- denoised$audio_latents diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md index fb53187..066630f 100644 --- a/inst/REFERENCES.md +++ b/inst/REFERENCES.md @@ -14,6 +14,10 @@ module; this file documents the actual lineage, idea by idea. | Distilled sigma schedules, default negative prompt | diffusers `pipelines/ltx2/utils.py` (Apache-2.0); values originate from Lightricks' LTX-2 release (numeric facts) | | Single-file checkpoint key layout, embedded config, `model_version` metadata | Format facts of the official Lightricks checkpoints, cross-checked against diffusers `scripts/convert_ltx2_to_diffusers.py` (Apache-2.0) | | Gemma3 text encoder | Ported from HuggingFace **transformers** (Apache-2.0) | +| LTX-2.5 transformer flags (`ff_bias = FALSE` video FFN, `keyframes_abs_pos_embedding` after `proj_in`) | diffusers `models/transformers/transformer_ltx2.py` (Apache-2.0); flag values from the embedded config of the official LTX-2.5 checkpoints (format facts) | +| Gemma 4 text encoder (LTX-2.5) | Ported from HuggingFace **transformers** (Apache-2.0) `models/gemma4_unified/modeling_gemma4_unified.py` (text stack) and `modeling_rope_utils.py` (`proportional` RoPE); the LTX encode recipe (strip, left pad, all 49 hidden states stacked) from diffusers `modular_pipelines/ltx2/encoders.py` (Apache-2.0). Parity fixtures: `tools/gen_fixtures_gemma4.py` | +| Gemma 4 tokenizer | Re-implements the HuggingFace **tokenizers** BPE pipeline described by the shipped `tokenizer.json` (normalizer, byte fallback, ranked merges, post-processor template); verified against `tokenizers` ids | +| Ancestral (SDE) Euler step for LTX-2.5, eta 1 | Ported from diffusers `schedulers/scheduling_ltx_euler_ancestral_rf.py` (Apache-2.0). That LTX-2.5+ checkpoints use it, and the stage noise seeds `seed + 10000` / `seed + 20000`, are numeric facts from Lightricks' LTX-2 release (`ANCESTRAL_SAMPLER_SINCE_VERSION`, LTX-2 Community License; no code taken) | ## FLUX.1-schnell diff --git a/inst/tinytest/fixtures/gemma4_real_tokenizer_cases.json b/inst/tinytest/fixtures/gemma4_real_tokenizer_cases.json new file mode 100644 index 0000000..30c52c3 --- /dev/null +++ b/inst/tinytest/fixtures/gemma4_real_tokenizer_cases.json @@ -0,0 +1 @@ +[{"text": "A golden retriever runs across a sunlit meadow, slow motion.", "ids": [236776, 13935, 102202, 8784, 3418, 496, 3768, 18114, 35280, 236764, 5111, 5776, 236761]}, {"text": "Close-up: rain drips off a neon sign that reads \"OPEN 24/7\" — cinematic.", "ids": [9564, 236772, 1048, 236787, 6927, 225274, 1135, 496, 23915, 1519, 600, 20838, 623, 38764, 236743, 236778, 236812, 236786, 236832, 236775, 2192, 60420, 236761]}, {"text": "一只猫在窗台上睡觉。", "ids": [118206, 239959, 237075, 239875, 211722, 163381, 236924]}, {"text": " two spaces and a tab\there ", "ids": [138, 13498, 138, 35220, 532, 496, 6937, 255968, 8472, 236743]}, {"text": "Emoji test 🚀🌙 and café crème brûlée", "ids": [127268, 1594, 236743, 242015, 244876, 532, 33443, 100479, 125501, 155666]}] \ No newline at end of file diff --git a/inst/tinytest/fixtures/gemma4_text.safetensors b/inst/tinytest/fixtures/gemma4_text.safetensors new file mode 100644 index 0000000..1b7c48b Binary files /dev/null and b/inst/tinytest/fixtures/gemma4_text.safetensors differ diff --git a/inst/tinytest/fixtures/gemma4_tiny_ckpt.safetensors b/inst/tinytest/fixtures/gemma4_tiny_ckpt.safetensors new file mode 100644 index 0000000..7b59fab Binary files /dev/null and b/inst/tinytest/fixtures/gemma4_tiny_ckpt.safetensors differ diff --git a/inst/tinytest/fixtures/gemma4_tiny_tokenizer.json b/inst/tinytest/fixtures/gemma4_tiny_tokenizer.json new file mode 100644 index 0000000..42c5488 --- /dev/null +++ b/inst/tinytest/fixtures/gemma4_tiny_tokenizer.json @@ -0,0 +1 @@ +{"version": "1.0", "truncation": null, "padding": null, "added_tokens": [{"id": 0, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 1, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 2, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 3, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}, {"id": 4, "content": "", "single_word": false, "lstrip": false, "rstrip": false, "normalized": false, "special": true}], "normalizer": {"type": "Replace", "pattern": {"String": " "}, "content": "▁"}, "pre_tokenizer": {"type": "Split", "pattern": {"String": " "}, "behavior": "MergedWithPrevious", "invert": false}, "post_processor": {"type": "TemplateProcessing", "single": [{"Sequence": {"id": "A", "type_id": 0}}], "pair": [{"Sequence": {"id": "A", "type_id": 0}}, {"Sequence": {"id": "B", "type_id": 1}}], "special_tokens": {}}, "decoder": null, "model": {"type": "BPE", "dropout": null, "unk_token": "", "continuing_subword_prefix": null, "end_of_word_suffix": null, "fuse_unk": true, "byte_fallback": true, "ignore_merges": false, "vocab": {"": 0, "": 1, "": 2, "": 3, "": 4, "<0x00>": 5, "<0x01>": 6, "<0x02>": 7, "<0x03>": 8, "<0x04>": 9, "<0x05>": 10, "<0x06>": 11, "<0x07>": 12, "<0x08>": 13, "<0x09>": 14, "<0x0A>": 15, "<0x0B>": 16, "<0x0C>": 17, "<0x0D>": 18, "<0x0E>": 19, "<0x0F>": 20, "<0x10>": 21, "<0x11>": 22, "<0x12>": 23, "<0x13>": 24, "<0x14>": 25, "<0x15>": 26, "<0x16>": 27, "<0x17>": 28, "<0x18>": 29, "<0x19>": 30, "<0x1A>": 31, "<0x1B>": 32, "<0x1C>": 33, "<0x1D>": 34, "<0x1E>": 35, "<0x1F>": 36, "<0x20>": 37, "<0x21>": 38, "<0x22>": 39, "<0x23>": 40, "<0x24>": 41, "<0x25>": 42, "<0x26>": 43, "<0x27>": 44, "<0x28>": 45, "<0x29>": 46, "<0x2A>": 47, "<0x2B>": 48, "<0x2C>": 49, "<0x2D>": 50, "<0x2E>": 51, "<0x2F>": 52, "<0x30>": 53, "<0x31>": 54, "<0x32>": 55, "<0x33>": 56, "<0x34>": 57, "<0x35>": 58, "<0x36>": 59, "<0x37>": 60, "<0x38>": 61, "<0x39>": 62, "<0x3A>": 63, "<0x3B>": 64, "<0x3C>": 65, "<0x3D>": 66, "<0x3E>": 67, "<0x3F>": 68, "<0x40>": 69, "<0x41>": 70, "<0x42>": 71, "<0x43>": 72, "<0x44>": 73, "<0x45>": 74, "<0x46>": 75, "<0x47>": 76, "<0x48>": 77, "<0x49>": 78, "<0x4A>": 79, "<0x4B>": 80, "<0x4C>": 81, "<0x4D>": 82, "<0x4E>": 83, "<0x4F>": 84, "<0x50>": 85, "<0x51>": 86, "<0x52>": 87, "<0x53>": 88, "<0x54>": 89, "<0x55>": 90, "<0x56>": 91, "<0x57>": 92, "<0x58>": 93, "<0x59>": 94, "<0x5A>": 95, "<0x5B>": 96, "<0x5C>": 97, "<0x5D>": 98, "<0x5E>": 99, "<0x5F>": 100, "<0x60>": 101, "<0x61>": 102, "<0x62>": 103, "<0x63>": 104, "<0x64>": 105, "<0x65>": 106, "<0x66>": 107, "<0x67>": 108, "<0x68>": 109, "<0x69>": 110, "<0x6A>": 111, "<0x6B>": 112, "<0x6C>": 113, "<0x6D>": 114, "<0x6E>": 115, "<0x6F>": 116, "<0x70>": 117, "<0x71>": 118, "<0x72>": 119, "<0x73>": 120, "<0x74>": 121, "<0x75>": 122, "<0x76>": 123, "<0x77>": 124, "<0x78>": 125, "<0x79>": 126, "<0x7A>": 127, "<0x7B>": 128, "<0x7C>": 129, "<0x7D>": 130, "<0x7E>": 131, "<0x7F>": 132, "<0x80>": 133, "<0x81>": 134, "<0x82>": 135, "<0x83>": 136, "<0x84>": 137, "<0x85>": 138, "<0x86>": 139, "<0x87>": 140, "<0x88>": 141, "<0x89>": 142, "<0x8A>": 143, "<0x8B>": 144, "<0x8C>": 145, "<0x8D>": 146, "<0x8E>": 147, "<0x8F>": 148, "<0x90>": 149, "<0x91>": 150, "<0x92>": 151, "<0x93>": 152, "<0x94>": 153, "<0x95>": 154, "<0x96>": 155, "<0x97>": 156, "<0x98>": 157, "<0x99>": 158, "<0x9A>": 159, "<0x9B>": 160, "<0x9C>": 161, "<0x9D>": 162, "<0x9E>": 163, "<0x9F>": 164, "<0xA0>": 165, "<0xA1>": 166, "<0xA2>": 167, "<0xA3>": 168, "<0xA4>": 169, "<0xA5>": 170, "<0xA6>": 171, "<0xA7>": 172, "<0xA8>": 173, "<0xA9>": 174, "<0xAA>": 175, "<0xAB>": 176, "<0xAC>": 177, "<0xAD>": 178, "<0xAE>": 179, "<0xAF>": 180, "<0xB0>": 181, "<0xB1>": 182, "<0xB2>": 183, "<0xB3>": 184, "<0xB4>": 185, "<0xB5>": 186, "<0xB6>": 187, "<0xB7>": 188, "<0xB8>": 189, "<0xB9>": 190, "<0xBA>": 191, "<0xBB>": 192, "<0xBC>": 193, "<0xBD>": 194, "<0xBE>": 195, "<0xBF>": 196, "<0xC0>": 197, "<0xC1>": 198, "<0xC2>": 199, "<0xC3>": 200, "<0xC4>": 201, "<0xC5>": 202, "<0xC6>": 203, "<0xC7>": 204, "<0xC8>": 205, "<0xC9>": 206, "<0xCA>": 207, "<0xCB>": 208, "<0xCC>": 209, "<0xCD>": 210, "<0xCE>": 211, "<0xCF>": 212, "<0xD0>": 213, "<0xD1>": 214, "<0xD2>": 215, "<0xD3>": 216, "<0xD4>": 217, "<0xD5>": 218, "<0xD6>": 219, "<0xD7>": 220, "<0xD8>": 221, "<0xD9>": 222, "<0xDA>": 223, "<0xDB>": 224, "<0xDC>": 225, "<0xDD>": 226, "<0xDE>": 227, "<0xDF>": 228, "<0xE0>": 229, "<0xE1>": 230, "<0xE2>": 231, "<0xE3>": 232, "<0xE4>": 233, "<0xE5>": 234, "<0xE6>": 235, "<0xE7>": 236, "<0xE8>": 237, "<0xE9>": 238, "<0xEA>": 239, "<0xEB>": 240, "<0xEC>": 241, "<0xED>": 242, "<0xEE>": 243, "<0xEF>": 244, "<0xF0>": 245, "<0xF1>": 246, "<0xF2>": 247, "<0xF3>": 248, "<0xF4>": 249, "<0xF5>": 250, "<0xF6>": 251, "<0xF7>": 252, "<0xF8>": 253, "<0xF9>": 254, "<0xFA>": 255, "<0xFB>": 256, "<0xFC>": 257, "<0xFD>": 258, "<0xFE>": 259, "<0xFF>": 260, "a": 261, "b": 262, "c": 263, "d": 264, "e": 265, "h": 266, "i": 267, "l": 268, "m": 269, "n": 270, "o": 271, "r": 272, "s": 273, "t": 274, "w": 275, ".": 276, "▁": 277, "▁t": 278, "he": 279, "▁the": 280, "ll": 281, "ell": 282, "ow": 283, "▁c": 284, "▁ca": 285, "at": 286, "▁cat": 287, "▁s": 288, "it": 289, "▁sit": 290, "▁o": 291, "n▁": 292}, "merges": [["▁", "t"], ["h", "e"], ["▁t", "he"], ["l", "l"], ["e", "ll"], ["o", "w"], ["▁", "c"], ["▁c", "a"], ["a", "t"], ["▁ca", "t"], ["▁", "s"], ["i", "t"], ["▁s", "it"], ["▁", "o"], ["n", "▁"]]}} \ No newline at end of file diff --git a/inst/tinytest/fixtures/gemma4_tokenizer_cases.json b/inst/tinytest/fixtures/gemma4_tokenizer_cases.json new file mode 100644 index 0000000..211f8fa --- /dev/null +++ b/inst/tinytest/fixtures/gemma4_tokenizer_cases.json @@ -0,0 +1 @@ +[{"text": "the cat sat", "ids": [274, 279, 287, 288, 286]}, {"text": "hello cat", "ids": [279, 281, 271, 287]}, {"text": " the cat ", "ids": [280, 277, 287, 277]}, {"text": "tell", "ids": [274, 282]}, {"text": "caté ok", "ids": [263, 286, 200, 174, 291, 112]}, {"text": "thecat", "ids": [274, 279, 1, 263, 286]}, {"text": "sit on the mat.", "ids": [273, 289, 291, 270, 280, 277, 269, 286, 276]}] \ No newline at end of file diff --git a/inst/tinytest/test_checkpoint_ltx25.R b/inst/tinytest/test_checkpoint_ltx25.R new file mode 100644 index 0000000..d72c32f --- /dev/null +++ b/inst/tinytest/test_checkpoint_ltx25.R @@ -0,0 +1,155 @@ +# ltx25_open_checkpoint(): the split LTX-2.5 files presented as one +# ltx23_checkpoint under the 2.3 key names. + +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) + +local({ + dir <- tempfile("ltx25-ckpt-") + dir.create(dir) + on.exit(unlink(dir, recursive = TRUE), add = TRUE) + path <- function(n) file.path(dir, n) + save <- function(tensors, name, version = "2.5.0", config = NULL) { + meta <- list(model_version = version) + if (!is.null(config)) { + meta$config <- as.character(jsonlite::toJSON(config, auto_unbox = TRUE)) + } + if (is.na(version)) meta$model_version <- NULL + safetensors::safe_save_file(tensors, path(name), metadata = meta) + } + t <- function(v) torch::torch_tensor(v) + + save(list("model.diffusion_model.patchify_proj.weight" = t(c(1, 2)), + "model.diffusion_model.video_embeddings_connector.learnable_registers" = t(3)), + "dit.safetensors", + config = list(transformer = list(ff_bias = FALSE), scheduler = list(a = 1))) + save(list("model.layers.0.mlp.up_proj.weight" = t(9), + "text_embedding_projection.video_aggregate_embed.weight" = t(c(4, 5))), + "te.safetensors", version = NA) + save(list("decoder.up_blocks.0.conv.weight" = t(6), + "per_channel_statistics.mean-of-means" = t(7)), + "vae.safetensors", config = list(vae = list(latent_channels = 128L))) + save(list("audio_vae.decoder.conv_in.weight" = t(8), + "vocoder.vocoder.conv_in.weight" = t(10)), + "audio.safetensors", + config = list(audio_vae = list(x = 1), vocoder = list(y = 2))) + + ck <- ltx25_open_checkpoint(path("dit.safetensors"), path("te.safetensors"), + path("vae.safetensors"), path("audio.safetensors")) + expect_inherits(ck, "ltx23_checkpoint") + expect_equal(ck$version, "2.5.0") + expect_null(ck$format) + + groups <- ltx23_split_keys(ck$keys) + expect_equal(groups$dit, "model.diffusion_model.patchify_proj.weight") + expect_equal(sort(groups$connectors), + sort(c("model.diffusion_model.video_embeddings_connector.learnable_registers", + "text_embedding_projection.video_aggregate_embed.weight"))) + expect_equal(sort(groups$vae), sort(c("vae.decoder.up_blocks.0.conv.weight", + "vae.per_channel_statistics.mean-of-means"))) + expect_equal(groups$audio_vae, "audio_vae.decoder.conv_in.weight") + expect_equal(groups$vocoder, "vocoder.vocoder.conv_in.weight") + expect_equal(length(groups$other), 0L) + # The encoder's own weights are not part of the pipeline checkpoint + expect_false(any(grepl("^model\\.layers", ck$keys))) + + # Each key reads from its own file + get <- function(k) as.numeric(ck$handle$get_tensor(k)) + expect_equal(get("text_embedding_projection.video_aggregate_embed.weight"), c(4, 5)) + expect_equal(get("vae.per_channel_statistics.mean-of-means"), 7) + expect_equal(get("vocoder.vocoder.conv_in.weight"), 10) + expect_equal(get("model.diffusion_model.patchify_proj.weight"), c(1, 2)) + expect_error(ck$handle$get_tensor("nope"), "not found") + # The VAE map strips the prefix the composite added + expect_equal(ltx23_map_vae_key("vae.per_channel_statistics.mean-of-means"), + "latents_mean") + + # Configs merge from their own files + expect_false(ck$config$transformer$ff_bias) + expect_equal(ck$config$vae$latent_channels, 128L) + expect_equal(ck$config$vocoder$y, 2L) + + # Audio is optional + ck2 <- ltx25_open_checkpoint(path("dit.safetensors"), path("te.safetensors"), + path("vae.safetensors")) + expect_equal(length(ltx23_split_keys(ck2$keys)$vocoder), 0L) + + # Wrong files are refused + save(list("model.diffusion_model.patchify_proj.weight" = t(1)), "dit23.safetensors", + version = "2.3.0") + expect_error(ltx25_open_checkpoint(path("dit23.safetensors"), path("te.safetensors"), + path("vae.safetensors")), "2\\.5") + save(list("decoder.transformer_blocks.0.w" = t(1)), "vaediff.safetensors") + expect_error(ltx25_open_checkpoint(path("dit.safetensors"), path("te.safetensors"), + path("vaediff.safetensors")), "convolutional") + expect_error(ltx25_open_checkpoint(path("dit.safetensors"), path("vae.safetensors"), + path("vae.safetensors")), "text_embedding_projection") + + # --- One self-contained NF4 artifact from the split files ----------------- + q_key <- "model.diffusion_model.transformer_blocks.0.attn1.to_q.weight" + torch::torch_manual_seed(3) + dq <- list(t(c(1, 2)), t(3), torch::torch_randn(8L, 16L)) + names(dq) <- c("model.diffusion_model.patchify_proj.weight", + "model.diffusion_model.video_embeddings_connector.learnable_registers", + q_key) + save(dq, "ditq.safetensors", + config = list(transformer = list(ff_bias = FALSE))) + art <- path("bundle") + m <- ltx23_quantize_nf4( + ltx25_open_checkpoint(path("ditq.safetensors"), path("te.safetensors"), + path("vae.safetensors"), path("audio.safetensors")), + art, verbose = FALSE) + expect_equal(m$model_version, "2.5.0") + expect_equal(m$nf4_cast, 1L) + expect_true(all(startsWith(m$shards, "ltx2.5-nf4-"))) + expect_equal(length(m$source), 4L) + ab <- ltx23_open_fp8_checkpoint(art) + expect_equal(ab$format, "nf4") + expect_true(all(c(paste0(q_key, "_absmax"), + "vae.per_channel_statistics.mean-of-means", + "text_embedding_projection.video_aggregate_embed.weight", + "vocoder.vocoder.conv_in.weight") %in% ab$keys)) + expect_false(any(grepl("^model\\.layers", ab$keys))) + expect_equal(ab$config$vae$latent_channels, 128L) + expect_false(ab$config$transformer$ff_bias) + expect_equal(as.numeric(ab$handle$get_tensor("vocoder.vocoder.conv_in.weight")), 10) + + # A transformer that is already NF4 is bundled as is, not re-quantized + dit_art <- path("dit-nf4") + ltx23_quantize_nf4(path("ditq.safetensors"), dit_art, verbose = FALSE) + art2 <- path("bundle2") + m2 <- ltx23_quantize_nf4( + ltx25_open_checkpoint(dit_art, path("te.safetensors"), path("vae.safetensors"), + path("audio.safetensors")), + art2, verbose = FALSE) + expect_equal(m2$nf4_cast, 1L) + b2 <- ltx23_open_fp8_checkpoint(art2) + expect_equal(sort(b2$keys), sort(ab$keys)) + expect_true(torch::torch_equal(b2$handle$get_tensor(q_key), + ab$handle$get_tensor(q_key))) + expect_true(torch::torch_equal(b2$handle$get_tensor(paste0(q_key, "_absmax")), + ab$handle$get_tensor(paste0(q_key, "_absmax")))) + + # The Gemma 4 artifact is told apart from a Gemma 3 one by its manifest + g4 <- path("g4"); dir.create(g4) + jsonlite::write_json(list(model = "gemma4"), file.path(g4, "manifest.json"), + auto_unbox = TRUE) + g3 <- path("g3"); dir.create(g3) + jsonlite::write_json(list(format = "nf4"), file.path(g3, "manifest.json"), + auto_unbox = TRUE) + expect_true(diffuseR:::.is_gemma4_artifact(g4)) + expect_false(diffuseR:::.is_gemma4_artifact(g3)) + expect_false(diffuseR:::.is_gemma4_artifact(path("te.safetensors"))) + expect_false(diffuseR:::.is_gemma4_artifact(NULL)) +}) + +# An encoder already on the CPU encodes there; staging or the GPU decides otherwise +m <- torch::nn_linear(2, 2) +expect_equal(diffuseR:::.ltx23_text_encode_device(m, "cuda"), "cpu") +expect_equal(diffuseR:::.ltx23_text_encode_device("path", "cuda"), "cuda") diff --git a/inst/tinytest/test_gemma4.R b/inst/tinytest/test_gemma4.R new file mode 100644 index 0000000..559aa3e --- /dev/null +++ b/inst/tinytest/test_gemma4.R @@ -0,0 +1,118 @@ +# Parity of the Gemma 4 (LTX-2.5) text encoder and tokenizer against +# HF transformers / tokenizers fixtures (tools/gen_fixtures_gemma4.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 <- function(name) { + p <- system.file("tinytest", "fixtures", name, package = "diffuseR") + if (p == "") p <- file.path("fixtures", name) + p +} +if (!file.exists(fixture("gemma4_text.safetensors"))) exit_file("gemma4 fixtures missing") +max_abs_diff <- function(a, b) as.numeric((a - b)$abs()$max()) + +# --- Tokenizer ------------------------------------------------------------------- + +tok <- gemma4_tokenizer(fixture("gemma4_tiny_tokenizer.json")) +expect_inherits(tok, "gemma4_tokenizer") +# Gemma 4's post-processor adds no special tokens (no BOS) +expect_equal(length(tok$template$prefix) + length(tok$template$suffix), 0L) +cases <- jsonlite::fromJSON(fixture("gemma4_tokenizer_cases.json"), + simplifyVector = FALSE) +for (cs in cases) { + got <- encode_gemma4(tok, cs$text, max_length = NULL)$input_ids[[1]] + expect_identical(as.integer(got), as.integer(unlist(cs$ids)), info = cs$text) +} +# Left padding with and truncation keeping the start +enc <- encode_gemma4(tok, c("the cat sat", "tell"), max_length = 4L) +expect_identical(enc$input_ids[2, ], c(0L, 0L, 274L, 282L)) +expect_identical(enc$attention_mask[2, ], c(0L, 0L, 1L, 1L)) +expect_identical(enc$input_ids[1, ], c(274L, 279L, 287L, 288L)) + +# The real LTX-2.5 tokenizer, when a copy is available locally +real_tok <- Sys.getenv("GEMMA4_TOKENIZER_JSON") +if (at_home() && nzchar(real_tok) && file.exists(real_tok)) { + rtok <- gemma4_tokenizer(real_tok) + rcases <- jsonlite::fromJSON(fixture("gemma4_real_tokenizer_cases.json"), + simplifyVector = FALSE) + for (cs in rcases) { + got <- encode_gemma4(rtok, cs$text, max_length = NULL)$input_ids[[1]] + expect_identical(as.integer(got), as.integer(unlist(cs$ids)), info = cs$text) + } +} + +# --- Model through the single-file loader --------------------------------------- + +model <- load_gemma4_text_encoder(fixture("gemma4_tiny_ckpt.safetensors"), + dtype = "float32", pin = FALSE, verbose = FALSE) +expect_inherits(model, "Gemma4TextModel") +blk_full <- model$layers[[3]]$self_attn +expect_true(blk_full$k_eq_v) +expect_null(blk_full$v_proj) +expect_equal(blk_full$head_dim, 16L) +expect_equal(blk_full$num_kv_heads, 1L) +expect_false(model$layers[[1]]$self_attn$k_eq_v) + +fx <- safetensors::safe_load_file(fixture("gemma4_text.safetensors"), + framework = "torch") +torch::with_no_grad({ + out <- model(fx$input_ids, attention_mask = fx$attention_mask, + output_hidden_states = TRUE) +}) +stacked <- torch::torch_stack(out$hidden_states, dim = -1L) +expect_equal(as.integer(stacked$shape), as.integer(fx$hidden_states$shape)) +# Real tokens are the contract (the connectors replace pad positions) +valid <- fx$attention_mask$to(dtype = torch::torch_bool()) +expect_true(max_abs_diff(stacked[valid], fx$hidden_states[valid]) < 1e-4) +expect_true(max_abs_diff(out$last_hidden_state[valid], + fx$last_hidden_state[valid]) < 1e-4) +expect_true(as.logical(stacked$isfinite()$all()$item())) + +# The full pipeline step: left-padded encode returns the same stack layout +enc <- encode_with_gemma4(c("the cat sat", "tell"), model, tok, + max_sequence_length = 6L, device = "cpu", + verbose = FALSE) +expect_equal(as.integer(enc$prompt_embeds$shape), c(2L, 6L, 32L, 7L)) +expect_identical(as.integer(as.array(enc$prompt_attention_mask[2, ])), + c(0L, 0L, 0L, 0L, 1L, 1L)) + +# --- Prompt stripping matches Python str.strip() -------------------------------- + +strip <- diffuseR:::.py_strip +expect_identical(strip(" tell \t\n"), "tell") +expect_identical(strip(c(" a b ", " x ")), c("a b", "x")) +expect_identical(strip("in side"), "in side") +enc_ws <- encode_with_gemma4(" tell ", model, tok, + max_sequence_length = 6L, device = "cpu", + verbose = FALSE) +enc_plain <- encode_with_gemma4("tell", model, tok, max_sequence_length = 6L, + device = "cpu", verbose = FALSE) +expect_true(torch::torch_equal(enc_ws$prompt_attention_mask, + enc_plain$prompt_attention_mask)) +expect_true(max_abs_diff(enc_ws$prompt_embeds, enc_plain$prompt_embeds) == 0) + +# --- Default (bfloat16) loader dtype -------------------------------------------- + +# The bf16 mask minimum must be finite in bf16 (float32's would be -Inf) +bmin <- diffuseR:::.gemma4_dtype_min(torch::torch_bfloat16()) +expect_true(is.finite(as.numeric(torch::torch_tensor(bmin, + dtype = torch::torch_bfloat16())$to(dtype = torch::torch_float32())))) + +model_bf <- load_gemma4_text_encoder(fixture("gemma4_tiny_ckpt.safetensors"), + pin = FALSE, verbose = FALSE) +expect_equal(model_bf$embed_tokens$weight$dtype$.type(), "BFloat16") +# "tell" in 6 slots leaves fully padded query rows: they must stay finite +enc_bf <- encode_with_gemma4(c("the cat sat", "tell"), model_bf, tok, + max_sequence_length = 6L, device = "cpu", + verbose = FALSE) +emb_bf <- enc_bf$prompt_embeds$to(dtype = torch::torch_float32()) +expect_true(as.logical(emb_bf$isfinite()$all()$item())) +ref <- enc$prompt_embeds +expect_true(max_abs_diff(emb_bf, ref) / as.numeric(ref$abs()$max()) < 0.05) diff --git a/inst/tinytest/test_gemma4_nf4.R b/inst/tinytest/test_gemma4_nf4.R new file mode 100644 index 0000000..aa1357b --- /dev/null +++ b/inst/tinytest/test_gemma4_nf4.R @@ -0,0 +1,158 @@ +# gemma4_quantize_nf4 + load_gemma4_nf4: round trip through the tiny +# LTX-2.5-layout checkpoint (tools/gen_fixtures_gemma4.py), compared with +# the dense load; the artifact keeps the connectors' projection and the +# tokenizer; a missing resident is a hard error. + +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 <- function(name) { + p <- system.file("tinytest", "fixtures", name, package = "diffuseR") + if (p == "") p <- file.path("fixtures", name) + p +} +src <- fixture("gemma4_tiny_ckpt.safetensors") +if (!file.exists(src)) exit_file("gemma4 fixtures missing") + +# Tokenizer straight from the encoder file's tokenizer_json tensor +tok_json <- gemma4_tokenizer(fixture("gemma4_tiny_tokenizer.json")) +tok_file <- gemma4_tokenizer(src) +expect_identical(encode_gemma4(tok_file, "the cat sat", NULL)$input_ids, + encode_gemma4(tok_json, "the cat sat", NULL)$input_ids) + +dense <- load_gemma4_text_encoder(src, dtype = "float32", pin = FALSE, + verbose = FALSE) +fx <- safetensors::safe_load_file(fixture("gemma4_text.safetensors"), + framework = "torch") +valid <- fx$attention_mask$to(dtype = torch::torch_bool()) + +art_seen <- local({ + art <- tempfile("gemma4-nf4-") + on.exit(unlink(art, recursive = TRUE), add = TRUE) + + manifest <- gemma4_quantize_nf4(src, art, verbose = FALSE) + expect_equal(manifest$model, "gemma4") + # 4 sliding layers x 7 projections + 2 full (k == v, no v_proj) x 6 + expect_equal(manifest$n_cast, 40L) + # Self-contained: tokenizer and the connectors' projection travel along + tok_art <- gemma4_tokenizer(art) + expect_identical(encode_gemma4(tok_art, "hello cat", NULL)$input_ids, + encode_gemma4(tok_json, "hello cat", NULL)$input_ids) + keys <- unlist(lapply(manifest$shards, function(s) { + names(safetensors::safe_load_file(file.path(art, s), framework = "torch")) + })) + expect_true("text_embedding_projection.video_aggregate_embed.weight" %in% keys) + expect_false(any(grepl("^vision_model|tokenizer_json", keys))) + # A second call reuses the artifact + expect_identical(gemma4_quantize_nf4(src, art, verbose = FALSE)$shards, + manifest$shards) + # ...but only a complete one: a missing tokenizer.json forces a rebuild + expect_equal(manifest$tokenizer, "tokenizer.json") + unlink(file.path(art, "tokenizer.json")) + expect_message(gemma4_quantize_nf4(src, art, verbose = TRUE), "Quantized") + expect_true(file.exists(file.path(art, "tokenizer.json"))) + + # Shards stay under shard_bytes; only a single oversized source tensor + # (NF4 packed + absmax pair, or one resident) may exceed it alone + local({ + small <- tempfile("gemma4-nf4-small-") + on.exit(unlink(small, recursive = TRUE), add = TRUE) + bound <- 20000 + m <- gemma4_quantize_nf4(src, small, shard_bytes = bound, verbose = FALSE) + expect_true(length(m$shards) > 1L) + for (s in m$shards) { + w <- safetensors::safe_load_file(file.path(small, s), framework = "torch") + payload <- sum(vapply(w, function(t) t$numel() * t$element_size(), 0)) + n_src <- length(setdiff(names(w), grep("_absmax$", names(w), value = TRUE))) + expect_true(payload <= bound * 1.01 || n_src == 1L, + info = sprintf("%s: %.0f bytes, %d tensors", s, payload, n_src)) + } + }) + + enc <- load_gemma4_text_encoder(art, device = "cpu", dtype = "float32", + pin = FALSE, verbose = FALSE) + expect_inherits(enc$layers[[1]]$self_attn$q_proj, "ltx23_nf4_linear") + expect_inherits(enc$layers[[3]]$mlp$down_proj, "ltx23_nf4_linear") + expect_null(enc$layers[[3]]$self_attn$v_proj) + # Residents equal the source at the artifact's storage dtype (bf16 when + # safetensors can write it), the layer_scalar buffers included + stored <- function(x) { + w <- safetensors::safe_load_file(file.path(art, manifest$shards[[1]]), + framework = "torch") + x$to(dtype = w[["embed_tokens.weight"]]$dtype)$to(dtype = torch::torch_float32()) + } + expect_true(torch::torch_equal(enc$embed_tokens$weight, + stored(dense$embed_tokens$weight))) + expect_true(torch::torch_equal(enc$layers[[2]]$layer_scalar, + stored(dense$layers[[2]]$layer_scalar))) + + # NF4 weight error per projection: 4-bit noise, not a layout bug + # (a transposed or misassigned weight would be ~1.4) + eye <- function(n) torch::torch_eye(n) + for (lp in list(c(1, "self_attn", "q_proj"), c(3, "self_attn", "k_proj"), + c(3, "mlp", "down_proj"))) { + l <- as.integer(lp[1]) + q <- enc$layers[[l]][[lp[2]]][[lp[3]]] + w <- dense$layers[[l]][[lp[2]]][[lp[3]]]$weight + deq <- torch::with_no_grad(q(eye(w$shape[2]))$t()) + rel_w <- as.numeric((deq - w)$norm() / w$norm()) + expect_true(rel_w < 0.15, info = sprintf("%s rel err %.3f", lp[3], rel_w)) + } + + # Load mechanics are exact: the NF4 model equals a dense twin carrying + # its dequantized weights and stored residents. (End-to-end NF4 noise on + # this random 6-layer model is large and says nothing about the 12B; + # real-model quality is checked on real weights.) + twin <- load_gemma4_text_encoder(src, dtype = "float32", pin = FALSE, + verbose = FALSE) + torch::with_no_grad({ + tp <- twin$named_parameters() + ep <- enc$named_parameters() + for (n in names(tp)) { + if (grepl("^layers\\.[0-9]+\\.(self_attn|mlp)\\.[a-z]+_proj\\.weight$", n)) { + segs <- strsplit(sub("\\.weight$", "", n), ".", fixed = TRUE)[[1]] + lin <- enc$layers[[as.integer(segs[2]) + 1L]][[segs[3]]][[segs[4]]] + tp[[n]]$copy_(lin(eye(tp[[n]]$shape[2]))$t()) + } else { + tp[[n]]$copy_(ep[[n]]) + } + } + for (n in names(twin$named_buffers())) { + twin$named_buffers()[[n]]$copy_(enc$named_buffers()[[n]]) + } + h_nf4 <- enc(fx$input_ids, attention_mask = fx$attention_mask)$last_hidden_state + h_twin <- twin(fx$input_ids, attention_mask = fx$attention_mask)$last_hidden_state + }) + expect_true(as.numeric((h_nf4[valid] - h_twin[valid])$abs()$max()) < 1e-5) + + # Default bf16 compute through the encode step stays finite + enc_bf <- load_gemma4_nf4(art, device = "cpu", pin = FALSE, verbose = FALSE) + e <- encode_with_gemma4(c("the cat sat", "tell"), enc_bf, tok_art, + max_sequence_length = 6L, device = "cpu", + verbose = FALSE) + expect_true(as.logical(e$prompt_embeds$to(dtype = torch::torch_float32())$ + isfinite()$all()$item())) + + # An artifact missing a resident is a hard error + local({ + art2 <- tempfile("gemma4-nf4-bad-") + on.exit(unlink(art2, recursive = TRUE), add = TRUE) + dir.create(art2) + for (s in manifest$shards) { + w <- safetensors::safe_load_file(file.path(art, s), framework = "torch") + w[["layers.1.layer_scalar"]] <- NULL + safetensors::safe_save_file(w, file.path(art2, s)) + } + file.copy(file.path(art, "manifest.json"), art2) + expect_error(load_gemma4_nf4(art2, device = "cpu", pin = FALSE, + verbose = FALSE), "not filled") + }) + art +}) +expect_false(dir.exists(art_seen)) diff --git a/inst/tinytest/test_jit_ltx23.R b/inst/tinytest/test_jit_ltx23.R index 78be617..7884291 100644 --- a/inst/tinytest/test_jit_ltx23.R +++ b/inst/tinytest/test_jit_ltx23.R @@ -16,13 +16,13 @@ dim <- 32L; heads <- 2L; head_dim <- 16L; cross_dim <- 32L adim <- 16L; aheads <- 2L; ahead_dim <- 8L; across_dim <- 16L B <- 2L; Sv <- 6L; Sa <- 5L; St <- 4L -make_block <- function() { +make_block <- function(...) { blk <- ltx23_transformer_block( dim = dim, num_attention_heads = heads, attention_head_dim = head_dim, cross_attention_dim = cross_dim, audio_dim = adim, audio_num_attention_heads = aheads, audio_attention_head_dim = ahead_dim, - audio_cross_attention_dim = across_dim + audio_cross_attention_dim = across_dim, ... ) blk$eval() @@ -127,6 +127,17 @@ torch::with_no_grad({ expect_true(max_abs_diff(out2[[1]], ref2[[1]]) < 1e-4) expect_true(max_abs_diff(out2[[2]], ref2[[2]]) < 1e-4) +# LTX-2.5 video FF has no biases: the stack takes the zero-length +# sentinel and must still match eager +blk25 <- make_block(ff_bias = FALSE) +expect_null(blk25$ff$net[[3]]$bias) +torch::with_no_grad({ + ref25 <- eager(blk25, h, ah) + out25 <- jit(list(blk25), h, ah) +}) +expect_true(max_abs_diff(out25[[1]], ref25[[1]]) < 1e-4) +expect_true(max_abs_diff(out25[[2]], ref25[[2]]) < 1e-4) + # The mask must actually mask: moving the -10000 changes the output enc_mask2 <- torch::torch_zeros(B, 1L, St) enc_mask2[, , 1L] <- -10000 diff --git a/inst/tinytest/test_ltx25_dit.R b/inst/tinytest/test_ltx25_dit.R new file mode 100644 index 0000000..0b0d4b4 --- /dev/null +++ b/inst/tinytest/test_ltx25_dit.R @@ -0,0 +1,166 @@ +# LTX-2.5 transformer deltas on the shared 2.3 module tree: the video +# feed-forward has no biases and a learned keyframe marker joins +# proj_in. Both flags come from the checkpoint's embedded config. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +# --- Config flags ---------------------------------------------------------------- + +expect_identical(ltx23_dit_config_args(NULL), list()) +# A 2.3 config carries none of the keys: the defaults apply +expect_identical(ltx23_dit_config_args(list(transformer = list(num_layers = 48L))), + list()) +cfg25 <- list(transformer = list(ff_bias = FALSE, + use_keyframes_abs_pos_embedding = TRUE)) +expect_identical(ltx23_dit_config_args(cfg25), + list(ff_bias = FALSE, use_keyframes_abs_pos_embedding = TRUE)) + +# --- Module shape ---------------------------------------------------------------- + +tiny_cfg <- list( + in_channels = 4L, out_channels = 4L, + num_attention_heads = 2L, attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 4L, audio_out_channels = 4L, + audio_num_attention_heads = 2L, audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 1L +) +torch::torch_manual_seed(3) +m23 <- do.call(ltx23_transformer, tiny_cfg) +m25 <- do.call(ltx23_transformer, c(tiny_cfg, ltx23_dit_config_args(cfg25))) +blk <- m25$transformer_blocks[[1]] +expect_null(blk$ff$net[[1]]$proj$bias) +expect_null(blk$ff$net[[3]]$bias) +expect_false(is.null(blk$audio_ff$net[[1]]$proj$bias)) +expect_null(m23$keyframes_abs_pos_embedding) +expect_equal(as.integer(m25$keyframes_abs_pos_embedding$shape), c(1L, 16L)) +# Exactly the two video FF biases and the marker differ, as in the real +# checkpoints (2 x 48 biases fewer, one marker more) +n23 <- names(m23$named_parameters()) +n25 <- names(m25$named_parameters()) +expect_equal(sort(setdiff(n23, n25)), + sort(c("transformer_blocks.0.ff.net.0.proj.bias", + "transformer_blocks.0.ff.net.2.bias"))) +expect_equal(setdiff(n25, n23), "keyframes_abs_pos_embedding") + +# --- Keyframe marker ------------------------------------------------------------- + +m25$eval() +run <- function(model, ...) { + torch::torch_manual_seed(5) + torch::with_no_grad({ + model( + hidden_states = torch::torch_randn(1L, 24L, 4L), + audio_hidden_states = torch::torch_randn(1L, 5L, 4L), + encoder_hidden_states = torch::torch_randn(1L, 7L, 16L), + audio_encoder_hidden_states = torch::torch_randn(1L, 7L, 8L), + timestep = torch::torch_tensor(700), + sigma = torch::torch_tensor(700), + num_frames = 2L, height = 3L, width = 4L, audio_num_frames = 5L, + use_cross_timestep = TRUE, ... + ) + }) +} +max_abs_diff <- function(a, b) as.numeric((a - b)$abs()$max()) + +torch::with_no_grad(m25$keyframes_abs_pos_embedding$fill_(0.5)) +first_frame <- torch::torch_zeros(1L, 24L, 1L) +first_frame[1, 1:12, 1] <- 1 +base <- run(m25) +expect_true(max_abs_diff(run(m25, video_keyframes_mask = first_frame)$sample, + base$sample) > 1e-4) +# An all-zero mask marks nothing +expect_true(max_abs_diff(run(m25, video_keyframes_mask = torch::torch_zeros(1L, 24L, 1L))$sample, + base$sample) < 1e-7) + +# --- 2.5 checkpoint through the NF4 and fp8 artifacts --------------------------- + +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} + +to_official <- function(name) { + name <- gsub("av_cross_attn_video_scale_shift", "av_ca_video_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_video_a2v_gate", "av_ca_a2v_gate_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_scale_shift", "av_ca_audio_scale_shift_adaln_single", name, fixed = TRUE) + name <- gsub("av_cross_attn_audio_v2a_gate", "av_ca_v2a_gate_adaln_single", name, fixed = TRUE) + name <- gsub("video_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_video", name, fixed = TRUE) + name <- gsub("audio_a2v_cross_attn_scale_shift_table", "scale_shift_table_a2v_ca_audio", name, fixed = TRUE) + name <- gsub("prompt_adaln", "prompt_adaln_single", name, fixed = TRUE) + name <- sub("^audio_time_embed\\.", "audio_adaln_single.", name) + name <- sub("^time_embed\\.", "adaln_single.", name) + name <- gsub("proj_in", "patchify_proj", name, fixed = TRUE) + name <- gsub("norm_q", "q_norm", name, fixed = TRUE) + name <- gsub("norm_k", "k_norm", name, fixed = TRUE) + paste0("model.diffusion_model.", name) +} + +params <- m25$named_parameters() + +# Every temp file and directory is created and removed inside its own +# local() scope, so a failing expectation or loader error still cleans up +src_seen <- local({ + tensors <- list() + for (name in names(params)) { + tensors[[to_official(name)]] <- params[[name]]$detach()$clone() + } + src <- tempfile(fileext = ".safetensors") + on.exit(unlink(src), add = TRUE) + safetensors::safe_save_file(tensors, src, metadata = list( + model_version = "2.5.0", + config = as.character(jsonlite::toJSON(cfg25, auto_unbox = TRUE)) + )) + + # The single-file opener accepts 2.5 and parses the flags back out + ckpt25 <- ltx23_open_checkpoint(src) + expect_equal(ckpt25$version, "2.5.0") + expect_identical(ltx23_dit_config_args(ckpt25$config), + list(ff_bias = FALSE, use_keyframes_abs_pos_embedding = TRUE)) + + # Dense (bf16/fp32) loader, the ltx23_load_pipeline() path. No ff_bias + # argument: it must arrive through the embedded config, and an + # incomplete load stops inside the loader. Loading at float32 makes + # the round trip exact. + d <- do.call(diffuseR:::.ltx23_load_transformer, + c(list(ckpt = ckpt25, dtype = torch::torch_float32(), + verbose = FALSE), tiny_cfg)) + expect_null(d$transformer_blocks[[1]]$ff$net[[3]]$bias) + loaded <- d$named_parameters() + expect_equal(sort(names(loaded)), sort(names(params))) + same <- vapply(names(params), function(n) { + as.logical(torch::torch_equal(loaded[[n]], params[[n]])) + }, logical(1)) + expect_true(all(same)) + + for (fmt in c("nf4", "fp8")) local({ + dir <- tempfile(paste0(fmt, "_")) + on.exit(unlink(dir, recursive = TRUE), add = TRUE) + quantize <- if (fmt == "nf4") ltx23_quantize_nf4 else ltx23_quantize_fp8 + load <- if (fmt == "nf4") ltx23_load_transformer_nf4 else ltx23_load_transformer_fp8 + manifest <- quantize(src, dir, verbose = FALSE) + expect_equal(manifest[[paste0(fmt, "_cast")]], + sum(ltx23_is_fp8_cast_key(names(params))), info = fmt) + + # The flags reach the quantized loaders through the manifest config + q <- do.call(load, c(list(ckpt = ltx23_open_fp8_checkpoint(dir), + device = "cpu", verbose = FALSE), tiny_cfg)) + qblk <- q$transformer_blocks[[1]] + expect_null(qblk$ff$net[[3]]$bias, info = fmt) + expect_false(is.null(qblk$audio_ff$net[[3]]$bias), info = fmt) + expect_equal(as.numeric(q$keyframes_abs_pos_embedding$float()$mean()), 0.5, + info = fmt) + + q$to(dtype = torch::torch_float32()) + out <- run(q, video_keyframes_mask = first_frame) + expect_true(as.logical(out$sample$isfinite()$all()$item()), info = fmt) + }) + src +}) + +# The scopes above removed their fixtures +expect_false(file.exists(src_seen)) diff --git a/inst/tinytest/test_ltx25_sampler.R b/inst/tinytest/test_ltx25_sampler.R new file mode 100644 index 0000000..fdbcf51 --- /dev/null +++ b/inst/tinytest/test_ltx25_sampler.R @@ -0,0 +1,98 @@ +# The LTX sampler step: plain Euler for 2.3, rectified-flow ancestral +# Euler (eta 1) for 2.5, and the loop's seeded, state-restoring noise. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) +step <- diffuseR:::.ltx23_step +max_abs_diff <- function(a, b) as.numeric((a - b)$abs()$max()) + +torch::torch_manual_seed(11) +x <- torch::torch_randn(2L, 7L, 4L) +v <- torch::torch_randn(2L, 7L, 4L) + +# eta = 0 is the 2.3 Euler step, whatever noise is passed +expect_true(max_abs_diff(step(x, v, 0.9, 0.725), x + v * (0.725 - 0.9)) < 1e-7) +expect_true(max_abs_diff(step(x, v, 0.9, 0.725, eta = 0, noise = torch::torch_randn(2L, 7L, 4L)), + x + v * (0.725 - 0.9)) < 1e-7) +# The last ancestral step lands on the denoised estimate x - sigma * v +expect_true(max_abs_diff(step(x, v, 0.421875, 0, eta = 1, noise = x), x - v * 0.421875) < 1e-6) + +# Marginal check: x_t = (1 - s) x0 + s eps with the exact velocity +# v = eps - x0 must step to (1 - s_next) x0 + N(0, s_next^2) +local({ + torch::torch_manual_seed(12) + n <- 400000L + x0 <- 0.8 + for (pair in list(c(1.0, 0.99375), c(0.909375, 0.725), c(0.725, 0.421875))) { + s <- pair[1]; sn <- pair[2] + eps <- torch::torch_randn(n) + xt <- eps * s + x0 * (1 - s) + out <- step(xt, eps - x0, s, sn, eta = 1, noise = torch::torch_randn(n)) + expect_equal(out$mean()$item(), (1 - sn) * x0, tolerance = 0.01, info = s) + expect_equal(out$std()$item(), sn, tolerance = 0.01, info = s) + } + # Euler on the same input keeps the original noise draw scaled to s_next + eps <- torch::torch_randn(n) + xt <- eps * 0.9 + x0 * 0.1 + expect_true(max_abs_diff(step(xt, eps - x0, 0.9, 0.725), + eps * 0.725 + x0 * 0.275) < 1e-5) +}) + +# --- Sampler selection ----------------------------------------------------------- + +eta_for <- diffuseR:::.ltx23_sampler_eta +expect_equal(eta_for("auto", "2.3.0"), 0) +expect_equal(eta_for("auto", "2.5.0"), 1) +expect_equal(eta_for("auto", "2.10.1"), 1) +expect_equal(eta_for("auto", NULL), 0) +expect_equal(eta_for("euler", "2.5.0"), 0) +expect_equal(eta_for("euler_ancestral", "2.3.0"), 1) + +# --- The denoise loop on a tiny transformer ------------------------------------ + +model <- ltx23_transformer( + in_channels = 4L, out_channels = 4L, + num_attention_heads = 2L, attention_head_dim = 8L, + cross_attention_dim = 16L, + audio_in_channels = 4L, audio_out_channels = 4L, + audio_num_attention_heads = 2L, audio_attention_head_dim = 4L, + audio_cross_attention_dim = 8L, + num_layers = 1L +) +model$eval() +torch::torch_manual_seed(13) +lat <- torch::torch_randn(1L, 24L, 4L) +alat <- torch::torch_randn(1L, 5L, 4L) +venc <- torch::torch_randn(1L, 6L, 16L) +aenc <- torch::torch_randn(1L, 6L, 8L) +tmask <- torch::torch_ones(1L, 6L) +denoise <- function(eta, noise_seed = NULL) { + diffuseR:::.ltx23_denoise( + model, lat, alat, c(1.0, 0.909375, 0.725, 0.0), + venc, aenc, tmask, + latent_frames = 2L, latent_height = 3L, latent_width = 4L, + audio_num_frames = 5L, frame_rate = 24, device = "cpu", + compute_dtype = torch::torch_float32(), verbose = FALSE, + eta = eta, noise_seed = noise_seed + ) +} + +a1 <- denoise(1, noise_seed = 10007) +a2 <- denoise(1, noise_seed = 10007) +a3 <- denoise(1, noise_seed = 10008) +e0 <- denoise(0) +expect_true(max_abs_diff(a1$latents, a2$latents) == 0) +expect_true(max_abs_diff(a1$audio_latents, a2$audio_latents) == 0) +expect_true(max_abs_diff(a1$latents, a3$latents) > 1e-3) +expect_true(max_abs_diff(a1$latents, e0$latents) > 1e-3) +expect_true(as.logical(a1$latents$isfinite()$all()$item())) + +# A seeded ancestral run restores the caller's RNG stream +torch::torch_manual_seed(99) +invisible(denoise(1, noise_seed = 10007)) +after <- torch::torch_randn(5L) +torch::torch_manual_seed(99) +expect_true(max_abs_diff(after, torch::torch_randn(5L)) == 0) diff --git a/man/encode_gemma4.Rd b/man/encode_gemma4.Rd new file mode 100644 index 0000000..4509299 --- /dev/null +++ b/man/encode_gemma4.Rd @@ -0,0 +1,25 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{encode_gemma4} +\alias{encode_gemma4} +\title{Encode prompts with the Gemma 4 tokenizer} +\usage{ +encode_gemma4(tokenizer, texts, max_length = 1024L) +} +\arguments{ +\item{tokenizer}{A \code{\link{gemma4_tokenizer}}.} + +\item{texts}{Character vector of prompts.} + +\item{max_length}{Integer. Fixed sequence length, or NULL for ragged +unpadded output.} +} +\value{ +List with \code{input_ids} and \code{attention_mask}: integer + matrices [length(texts), max_length] (ragged lists when + \code{max_length} is NULL). Ids are 0-based. +} +\description{ +Tokenizes, applies the post-processor's special tokens (none for +Gemma 4), truncates to \code{max_length} keeping the start, and pads +on the left with \code{}, as the LTX pipelines do. +} diff --git a/man/encode_with_gemma4.Rd b/man/encode_with_gemma4.Rd new file mode 100644 index 0000000..1536413 --- /dev/null +++ b/man/encode_with_gemma4.Rd @@ -0,0 +1,41 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{encode_with_gemma4} +\alias{encode_with_gemma4} +\title{Encode prompts with Gemma 4 for LTX-2.5} +\usage{ +encode_with_gemma4( + prompts, + model, + tokenizer, + max_sequence_length = 1024L, + device = "cuda", + verbose = TRUE +) +} +\arguments{ +\item{prompts}{Character vector.} + +\item{model}{A \code{\link{gemma4_text_model}}. A pinned CPU-resident +model moves to \code{device} for the encode and back afterwards.} + +\item{tokenizer}{A \code{\link{gemma4_tokenizer}}.} + +\item{max_sequence_length}{Integer. Padded token length.} + +\item{device}{Character. Compute device.} + +\item{verbose}{Logical.} +} +\value{ +List with \code{prompt_embeds} \code{[batch, seq_len, + hidden_size, num_layers + 1]} and \code{prompt_attention_mask} + \code{[batch, seq_len]}, the same layout as + \code{\link{encode_with_gemma3}}. +} +\description{ +The LTX-2.5 text-encoder step: strips each prompt (Python +\code{str.strip()} semantics, Unicode spaces included), tokenizes with left +padding to \code{max_sequence_length}, and returns every hidden state +(embedding output plus each layer, the last final-normed) stacked on a +new last dimension for the connectors. +} diff --git a/man/gemma4_attention.Rd b/man/gemma4_attention.Rd new file mode 100644 index 0000000..4c5fce6 --- /dev/null +++ b/man/gemma4_attention.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_attention} +\alias{gemma4_attention} +\title{Gemma 4 text attention} +\usage{ +gemma4_attention(config, layer_idx) +} +\arguments{ +\item{config}{Model configuration (see \code{\link{load_gemma4_text_encoder}}).} + +\item{layer_idx}{Integer. 0-based layer index.} +} +\value{ +Module whose forward(hidden_states, attention_mask, cos, sin) + returns the attention output [B, S, hidden_size]. +} +\description{ +Gemma 4 text attention +} +\keyword{internal} diff --git a/man/gemma4_decoder_layer.Rd b/man/gemma4_decoder_layer.Rd new file mode 100644 index 0000000..8bcc045 --- /dev/null +++ b/man/gemma4_decoder_layer.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_decoder_layer} +\alias{gemma4_decoder_layer} +\title{Gemma 4 decoder layer} +\usage{ +gemma4_decoder_layer(config, layer_idx) +} +\arguments{ +\item{config}{Model configuration.} + +\item{layer_idx}{Integer. 0-based layer index.} +} +\value{ +Module whose forward(hidden_states, attention_mask, cos, sin) + returns the layer output, scaled by the stored \code{layer_scalar}. +} +\description{ +Gemma 4 decoder layer +} +\keyword{internal} diff --git a/man/gemma4_quantize_nf4.Rd b/man/gemma4_quantize_nf4.Rd new file mode 100644 index 0000000..3ef285e --- /dev/null +++ b/man/gemma4_quantize_nf4.Rd @@ -0,0 +1,52 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_quantize_nf4} +\alias{gemma4_quantize_nf4} +\title{Quantize the LTX-2.5 Gemma 4 text encoder to NF4 shards} +\usage{ +gemma4_quantize_nf4( + path, + output_dir = NULL, + shard_bytes = 1.9e+09, + force = FALSE, + verbose = TRUE +) +} +\arguments{ +\item{path}{The LTX-2.5 text encoder file +(\code{gemma4-12b-with-proj-ltx-2.5-bf16.safetensors}).} + +\item{output_dir}{Output directory (default: \code{gemma4-nf4} under +\code{tools::R_user_dir}).} + +\item{shard_bytes}{Numeric. Target shard size in bytes; the 1.9e9 +default keeps shards readable by safetensors older than 0.3.0.} + +\item{force}{Logical. Re-quantize even if a valid manifest exists.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, the manifest list. +} +\description{ +Streams the LTX-2.5 encoder file tensor by tensor. The per-layer +projection weights (q/k/v/o and gate/up/down; 328 tensors in the 12B, +since full-attention layers have no v_proj) are stored as NF4 (packed +uint8 + \code{_absmax} float32 blocks); embeddings, norms and the +per-layer \code{layer_scalar} are copied at the resident dtype. Vision +and audio towers are dropped. +} +\details{ +The artifact is self-contained for the LTX-2.5 text side: it also +keeps the connectors' \code{text_embedding_projection.*} weights +(resident dtype, original keys) and writes the embedded tokenizer as +\code{tokenizer.json}, so the 26 GB source file is not needed after +quantizing. Load it with \code{\link{load_gemma4_nf4}} (or +\code{\link{load_gemma4_text_encoder}} on the directory) and the +tokenizer with \code{gemma4_tokenizer(output_dir)}. + +Peak host memory is about one shard (\code{shard_bytes}) plus the +largest single source tensor and its save buffer (the 2 GB bf16 +embedding); measured under a 20 GB cgroup cap on the 12B. + +} diff --git a/man/gemma4_rms_norm.Rd b/man/gemma4_rms_norm.Rd new file mode 100644 index 0000000..e20e65a --- /dev/null +++ b/man/gemma4_rms_norm.Rd @@ -0,0 +1,23 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_rms_norm} +\alias{gemma4_rms_norm} +\title{Gemma 4 RMS normalization} +\usage{ +gemma4_rms_norm(dim, eps = 1e-06, with_scale = TRUE) +} +\arguments{ +\item{dim}{Integer. Normalized dimension.} + +\item{eps}{Numeric. Epsilon.} + +\item{with_scale}{Logical. Carry a learned scale (FALSE for the value +norm).} +} +\value{ +Module whose forward(x) normalizes over the last dimension in + float32 and multiplies by the scale, returning \code{x}'s dtype. +} +\description{ +Gemma 4 RMS normalization +} +\keyword{internal} diff --git a/man/gemma4_text_model.Rd b/man/gemma4_text_model.Rd new file mode 100644 index 0000000..fc5aaaf --- /dev/null +++ b/man/gemma4_text_model.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_text_model} +\alias{gemma4_text_model} +\title{Gemma 4 text model} +\usage{ +gemma4_text_model(config) +} +\arguments{ +\item{config}{Model configuration list: the checkpoint's +\code{text_config} fields (see \code{\link{load_gemma4_text_encoder}}).} +} +\value{ +Module whose forward(input_ids, attention_mask, + output_hidden_states) returns \code{list(last_hidden_state, + hidden_states)}. \code{input_ids} are 0-based; \code{hidden_states} + follows transformers: the state entering each layer, then the + final-normed output (\code{num_hidden_layers + 1} entries). +} +\description{ +The text stack of Gemma 4 Unified, as the LTX-2.5 text encoder. +} diff --git a/man/gemma4_tokenizer.Rd b/man/gemma4_tokenizer.Rd new file mode 100644 index 0000000..8e07872 --- /dev/null +++ b/man/gemma4_tokenizer.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{gemma4_tokenizer} +\alias{gemma4_tokenizer} +\title{Load the Gemma 4 tokenizer} +\usage{ +gemma4_tokenizer(source) +} +\arguments{ +\item{source}{Path to a \code{tokenizer.json}, a directory holding +one, or an LTX-2.5 text encoder \code{.safetensors} file, which +carries the same JSON in its \code{tokenizer_json} tensor.} +} +\value{ +A \code{gemma4_tokenizer}. +} +\description{ +Load the Gemma 4 tokenizer +} diff --git a/man/load_gemma4_nf4.Rd b/man/load_gemma4_nf4.Rd new file mode 100644 index 0000000..abb90fa --- /dev/null +++ b/man/load_gemma4_nf4.Rd @@ -0,0 +1,38 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_gemma4_nf4} +\alias{load_gemma4_nf4} +\title{Load the LTX-2.5 Gemma 4 text encoder from an NF4 artifact} +\usage{ +load_gemma4_nf4( + artifact_dir, + device = "cuda", + dtype = "bfloat16", + pin = getOption("diffuseR.pin_staging", TRUE), + verbose = TRUE +) +} +\arguments{ +\item{artifact_dir}{Directory produced by \code{\link{gemma4_quantize_nf4}}.} + +\item{device}{"cuda" or "cpu".} + +\item{dtype}{Compute dtype ("bfloat16" default, or "float32").} + +\item{pin}{Logical. When loading to the CPU, page-lock the weights so +\code{\link{encode_with_gemma4}} can move the model to the GPU per +encode at DMA speed. Default follows +\code{options(diffuseR.pin_staging)}.} + +\item{verbose}{Logical.} +} +\value{ +A \code{gemma4_text_model}. +} +\description{ +Builds the model as a skeleton at the compute dtype, swaps the +projection linears for NF4 modules filled from the artifact +(dequantized per forward), copies the residents (parameters and the +\code{layer_scalar} buffers), and hard-errors on anything the artifact +does not fill. The connectors' projection weights in the artifact are +left for the connector loader. +} diff --git a/man/load_gemma4_text_encoder.Rd b/man/load_gemma4_text_encoder.Rd new file mode 100644 index 0000000..2389aa2 --- /dev/null +++ b/man/load_gemma4_text_encoder.Rd @@ -0,0 +1,40 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_gemma4_text_encoder} +\alias{load_gemma4_text_encoder} +\title{Load the Gemma 4 text encoder (LTX-2.5)} +\usage{ +load_gemma4_text_encoder( + path, + device = "cpu", + dtype = "bfloat16", + pin = getOption("diffuseR.pin_staging", TRUE), + verbose = TRUE +) +} +\arguments{ +\item{path}{Path to the LTX-2.5 text encoder file (e.g. +\code{gemma4-12b-with-proj-ltx-2.5-bf16.safetensors}), which embeds +its config as \code{gemma_config} metadata, or an NF4 artifact +directory from \code{\link{gemma4_quantize_nf4}} (dispatched to +\code{\link{load_gemma4_nf4}}).} + +\item{device}{Character. Device for the loaded model.} + +\item{dtype}{Character. "bfloat16" (checkpoint native) or "float32".} + +\item{pin}{Logical. When loading to the CPU, page-lock the weights so +\code{\link{encode_with_gemma4}} can move the model to the GPU per +encode at DMA speed. Default follows +\code{options(diffuseR.pin_staging)}.} + +\item{verbose}{Logical.} +} +\value{ +A \code{gemma4_text_model}. +} +\description{ +Streams the text weights of an LTX-2.5 Gemma 4 checkpoint into a +\code{\link{gemma4_text_model}} skeleton, one tensor at a time, so +the 26 GB bf16 file is never fully in memory. The vision and audio +towers in the same file are skipped. +} diff --git a/man/ltx23_dit_config_args.Rd b/man/ltx23_dit_config_args.Rd new file mode 100644 index 0000000..ac2f874 --- /dev/null +++ b/man/ltx23_dit_config_args.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx23_dit_config_args} +\alias{ltx23_dit_config_args} +\title{Transformer arguments from a checkpoint's embedded config} +\usage{ +ltx23_dit_config_args(config) +} +\arguments{ +\item{config}{The parsed embedded config (\code{ckpt$config}), or NULL.} +} +\value{ +Named list of \code{\link{ltx23_transformer}} arguments. +} +\description{ +Reads the version-dependent flags out of the \code{transformer} entry +of an LTX single-file checkpoint's embedded config. Every other +dimension is shared by the 2.3 and 2.5 22B checkpoints and stays at +the \code{\link{ltx23_transformer}} defaults. Keys absent from the +config (any 2.3 checkpoint) are left out, so the defaults apply. +} diff --git a/man/ltx23_feed_forward.Rd b/man/ltx23_feed_forward.Rd index aaa197f..b99c329 100644 --- a/man/ltx23_feed_forward.Rd +++ b/man/ltx23_feed_forward.Rd @@ -3,12 +3,15 @@ \alias{ltx23_feed_forward} \title{LTX feed-forward layer} \usage{ -ltx23_feed_forward(dim, mult = 4L) +ltx23_feed_forward(dim, mult = 4L, bias = TRUE) } \arguments{ \item{dim}{Integer. Input/output dimension.} \item{mult}{Integer. Hidden dimension multiplier.} + +\item{bias}{Logical. Both linears carry a bias (FALSE for the LTX-2.5 +video stream).} } \value{ Module whose forward(x) returns the projected states passed diff --git a/man/ltx23_load_pipeline.Rd b/man/ltx23_load_pipeline.Rd index e86f851..c982fdb 100644 --- a/man/ltx23_load_pipeline.Rd +++ b/man/ltx23_load_pipeline.Rd @@ -20,7 +20,9 @@ ltx23_load_pipeline( \code{ltx-2.3-22b-distilled-1.1.safetensors}) or to an fp8 artifact directory produced by \code{\link{ltx23_quantize_fp8}}. With the fp8 artifact, the transformer loads with CPU-resident fp8 weights that -stream to \code{device} during the forward pass.} +stream to \code{device} during the forward pass. An open +\code{ltx23_checkpoint} is also accepted, such as the LTX-2.5 +component set from \code{\link{ltx25_open_checkpoint}}.} \item{device}{Character. Device for the small components (VAEs, vocoder, connectors) and, with fp8, the transformer residents.} diff --git a/man/ltx23_open_checkpoint.Rd b/man/ltx23_open_checkpoint.Rd index 980ada2..80d0abc 100644 --- a/man/ltx23_open_checkpoint.Rd +++ b/man/ltx23_open_checkpoint.Rd @@ -3,13 +3,14 @@ \alias{ltx23_open_checkpoint} \title{Open an LTX-2.3 checkpoint} \usage{ -ltx23_open_checkpoint(path, require_version = "2.3") +ltx23_open_checkpoint(path, require_version = c("2.3", "2.5")) } \arguments{ \item{path}{Path to the checkpoint .safetensors file.} -\item{require_version}{Character. Required \code{model_version} prefix -(default "2.3"). Set to NULL to skip the check.} +\item{require_version}{Character vector. Accepted \code{model_version} +prefixes (default \code{c("2.3", "2.5")}, the 22B generations this +module tree loads). Set to NULL to skip the check.} } \value{ An object of class \code{ltx23_checkpoint}: a list with diff --git a/man/ltx23_quantize_nf4.Rd b/man/ltx23_quantize_nf4.Rd index 9b1566d..a6a4edf 100644 --- a/man/ltx23_quantize_nf4.Rd +++ b/man/ltx23_quantize_nf4.Rd @@ -12,7 +12,12 @@ ltx23_quantize_nf4( ) } \arguments{ -\item{checkpoint_path}{Source .safetensors (bf16 single file).} +\item{checkpoint_path}{Source .safetensors (bf16 single file), or an +open \code{ltx23_checkpoint}. For LTX-2.5, pass +\code{\link{ltx25_open_checkpoint}} to write one self-contained +artifact (transformer, connectors and text projection, VAEs, +vocoder); when its transformer is already an NF4 artifact, the +NF4 tensors are bundled as they are, with no re-quantization.} \item{output_dir}{Output directory for shards + manifest; NULL (the default) resolves under tools::R_user_dir("diffuseR", "data").} diff --git a/man/ltx23_transformer.Rd b/man/ltx23_transformer.Rd index 4e2cf25..717abb1 100644 --- a/man/ltx23_transformer.Rd +++ b/man/ltx23_transformer.Rd @@ -38,7 +38,10 @@ ltx23_transformer( timestep_scale_multiplier = 1000, cross_attn_timestep_scale_multiplier = 1000, rope_type = "split", - perturbed_attn = TRUE + perturbed_attn = TRUE, + ff_bias = TRUE, + audio_ff_bias = TRUE, + use_keyframes_abs_pos_embedding = FALSE ) } \arguments{ @@ -54,6 +57,11 @@ ltx23_transformer( \item{rope_type}{"split" (LTX-2.3) or "interleaved".} +\item{use_keyframes_abs_pos_embedding}{Logical. Carry the LTX-2.5 +learned keyframe marker, added after \code{proj_in} to the tokens +flagged by the forward's \code{video_keyframes_mask} +(\code{[B, S, 1]}, non-zero on keyframe tokens).} + \item{in_channels,out_channels}{Integers. Video latent channels.} \item{patch_size,patch_size_t}{Integers. Video patch sizes.} @@ -76,6 +84,9 @@ ltx23_transformer( the a2v/v2a gates).} \item{gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn}{LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints).} + +\item{ff_bias,audio_ff_bias}{Logicals. Feed-forward biases per stream +(\code{ff_bias = FALSE} for LTX-2.5).} } \value{ Module whose forward(hidden_states, ...) returns @@ -88,3 +99,7 @@ Dual-stream audio/video DiT. Text embeddings arrive already projected to the video (\code{inner_dim}) and audio (\code{audio_inner_dim}) dimensions by the connector modules. } +\seealso{ +\code{\link{ltx23_dit_config_args}} to derive these from a + checkpoint's embedded config. +} diff --git a/man/ltx23_transformer_block.Rd b/man/ltx23_transformer_block.Rd index ae190aa..32adb7a 100644 --- a/man/ltx23_transformer_block.Rd +++ b/man/ltx23_transformer_block.Rd @@ -19,7 +19,9 @@ ltx23_transformer_block( eps = 1e-06, elementwise_affine = FALSE, rope_type = "split", - perturbed_attn = TRUE + perturbed_attn = TRUE, + ff_bias = TRUE, + audio_ff_bias = TRUE ) } \arguments{ @@ -45,6 +47,9 @@ ltx23_transformer_block( \item{video_cross_attn_adaln,audio_cross_attn_adaln}{Logicals. LTX-2.3 text cross-attention modulation (9 mod params instead of 6).} + +\item{ff_bias,audio_ff_bias}{Logicals. Feed-forward biases per stream +(LTX-2.5 drops the video ones).} } \value{ Module whose forward(hidden_states, ...) returns diff --git a/man/ltx25_open_checkpoint.Rd b/man/ltx25_open_checkpoint.Rd new file mode 100644 index 0000000..56b17e5 --- /dev/null +++ b/man/ltx25_open_checkpoint.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{ltx25_open_checkpoint} +\alias{ltx25_open_checkpoint} +\title{Open the LTX-2.5 component files as one checkpoint} +\usage{ +ltx25_open_checkpoint(transformer, text_encoder, video_vae, audio_vae = NULL) +} +\arguments{ +\item{transformer}{The 2.5 transformer file +(\code{ltx-2.5-22b-distilled-transformer-bf16.safetensors}) or an +fp8/NF4 artifact directory quantized from it.} + +\item{text_encoder}{The 2.5 Gemma 4 text encoder file, or its NF4 +artifact directory (\code{\link{gemma4_quantize_nf4}} keeps the +projection); only \code{text_embedding_projection.*} is used here.} + +\item{video_vae}{The 2.5 video VAE file. Use the convolutional decoder +(\code{ltx-2.5-video-vae-conv-bf16.safetensors}); the default +diffusion decoder is not supported.} + +\item{audio_vae}{Optional 2.5 audio VAE + vocoder file +(\code{ltx-2.5-audio-vae-bf16.safetensors}).} +} +\value{ +An \code{ltx23_checkpoint} whose \code{handle$get_tensor()} + reads each key from its source file. +} +\description{ +LTX-2.3 shipped every component in one file; LTX-2.5 splits them: the +transformer (with the text connectors), the Gemma 4 text encoder file +(which carries the connectors' \code{text_embedding_projection}), the +video VAE, and the audio VAE plus vocoder. This presents them as a +single \code{ltx23_checkpoint} under the 2.3 key names (the video VAE +keys gain their \code{vae.} prefix), so \code{\link{ltx23_load_pipeline}} +and its component loaders work unchanged. +} diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 808e95c..ed2313f 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -27,7 +27,9 @@ resident_load( \item{verbose}{Print progress messages.} \item{text_encoder,tokenizer}{\code{ltx} only: paths to the Gemma3 encoder -artifact and the tokenizer directory. \code{\link{ltx23_load_pipeline}} +artifact and the tokenizer directory, or for LTX-2.5 the Gemma 4 NF4 +artifact (\code{\link{gemma4_quantize_nf4}}) for both, since it carries +its tokenizer. \code{\link{ltx23_load_pipeline}} does not load these -- \code{\link{txt2vid_ltx2}} takes them per call -- so a handle built without them can be activated and cannot generate. Given here they are loaded ONCE, pinned on the host, and passed to every diff --git a/man/tokenizer_gemma4.Rd b/man/tokenizer_gemma4.Rd new file mode 100644 index 0000000..326840f --- /dev/null +++ b/man/tokenizer_gemma4.Rd @@ -0,0 +1,14 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{tokenizer_gemma4} +\alias{tokenizer_gemma4} +\title{Gemma 4 SentencePiece-style BPE tokenizer} +\description{ +Native R port of the HuggingFace \code{tokenizers} pipeline in the +Gemma 4 \code{tokenizer.json} shipped with LTX-2.5: added-token +literals split out first, the \code{Replace " " -> U+2581} normalizer, +characters mapped to vocab ids with \code{<0xHH>} byte fallback, then +rank-ordered BPE merges over the whole string (the \code{Split} on +" " pre-tokenizer never matches after the normalizer). The +post-processor template decides special tokens; for Gemma 4 it adds +none, so there is no BOS. +} diff --git a/man/txt2vid_ltx2.Rd b/man/txt2vid_ltx2.Rd index ff5f978..d7953b7 100644 --- a/man/txt2vid_ltx2.Rd +++ b/man/txt2vid_ltx2.Rd @@ -36,7 +36,8 @@ txt2vid_ltx2( resident = character(), trim_frames = 0L, audio = NULL, - verbose = TRUE + verbose = TRUE, + sampler = c("auto", "euler", "euler_ancestral") ) } \arguments{ @@ -147,10 +148,18 @@ lines), FALSE = "silent". "progress" prints a one-line generation summary plus a denoise progress bar (interactive) or periodic step ticks (captured logs).} +\item{sampler}{One of "auto", "euler" or "euler_ancestral". "auto" +picks the sampler the checkpoint was distilled for: plain Euler +for LTX-2.3, ancestral Euler (eta 1) for LTX-2.5 and later. The +ancestral steps draw their noise from \code{seed + 10000} (stage 1) +and \code{seed + 20000} (stage 2), so the initial noise matches the +Euler run at the same seed.} + \item{text_encoder,tokenizer}{Gemma3 model and tokenizer (or paths; see \code{\link{load_gemma3_text_encoder}} and -\code{\link{gemma3_tokenizer}}). Ignored when \code{prompt_embeds} -is supplied.} +\code{\link{gemma3_tokenizer}}), or for LTX-2.5 the Gemma 4 model +and \code{\link{gemma4_tokenizer}} (selected by the tokenizer's +class). Ignored when \code{prompt_embeds} is supplied.} \item{width,height}{Integers. Output resolution (multiples of 32).} diff --git a/tools/gen_fixtures_gemma4.py b/tools/gen_fixtures_gemma4.py new file mode 100644 index 0000000..71e99f6 --- /dev/null +++ b/tools/gen_fixtures_gemma4.py @@ -0,0 +1,189 @@ +# Generate Gemma 4 (LTX-2.5 text encoder) parity fixtures for the R port. +# +# Uses HF transformers' Gemma4UnifiedTextModel (Apache-2.0) with a tiny +# random-init config that keeps every LTX-2.5 feature: 5:1-style sliding +# and full layers, full layers with a doubled head_dim, one KV head and +# k == v, proportional RoPE with partial_rotary_factor 0.25, v_norm and +# per-layer layer_scalar. All weights are perturbed off their init so +# norm scales (x * w, not x * (1 + w)) and layer scalars are exercised. +# +# Also builds a tiny tokenizer.json with the Gemma 4 pipeline (Replace +# " " -> U+2581, byte fallback, no BOS template) and records HF tokenizers +# ids for a set of prompts. Run via tools/gen_fixtures.sh (needs the +# tokenizers wheel, which transformers depends on). + +import json +import os + +import torch +from safetensors.torch import save_file +from tokenizers import Tokenizer +from transformers import Gemma4UnifiedTextConfig, Gemma4UnifiedTextModel + +OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures") +os.makedirs(OUT_DIR, exist_ok=True) + +torch.manual_seed(54) + +layer_types = ["sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "full_attention"] +text_config = dict( + vocab_size=300, # covers the tiny tokenizer below, for the encode test + hidden_size=32, + intermediate_size=64, + num_hidden_layers=len(layer_types), + num_attention_heads=4, + num_key_value_heads=2, + head_dim=8, + global_head_dim=16, + num_global_key_value_heads=1, + attention_k_eq_v=True, + layer_types=layer_types, + sliding_window=3, + rope_parameters={ + "full_attention": {"rope_type": "proportional", + "partial_rotary_factor": 0.25, + "rope_theta": 1_000_000.0}, + "sliding_attention": {"rope_type": "default", "rope_theta": 10_000.0}, + }, + hidden_activation="gelu_pytorch_tanh", + rms_norm_eps=1e-6, + max_position_embeddings=512, + num_kv_shared_layers=0, + hidden_size_per_layer_input=0, + enable_moe_block=False, + use_double_wide_mlp=False, + pad_token_id=0, +) +config = Gemma4UnifiedTextConfig(**text_config) +config._attn_implementation = "eager" +model = Gemma4UnifiedTextModel(config) +model.eval() + +with torch.no_grad(): + for name, p in model.named_parameters(): + if name.endswith("norm.weight"): + p.copy_(1.0 + 0.3 * torch.randn_like(p)) + else: + p.add_(0.05 * torch.randn_like(p)) + for layer in model.layers: + layer.layer_scalar.copy_(0.5 + torch.rand_like(layer.layer_scalar)) + +# Left-padded batch (LTX pads on the left), longer than the sliding +# window so the window mask binds +input_ids = torch.tensor([ + [5, 23, 61, 7, 19, 88, 42, 3, 9], + [0, 0, 0, 33, 14, 2, 71, 55, 12], +]) +attention_mask = (input_ids != 0).long() +attention_mask[0] = 1 + +with torch.no_grad(): + out = model(input_ids=input_ids, attention_mask=attention_mask, + output_hidden_states=True, use_cache=False) +hidden = torch.stack(out.hidden_states, dim=-1) # [B, S, H, L + 1] + +fx = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "hidden_states": hidden, + "last_hidden_state": out.last_hidden_state, +} +fx = {k: v.contiguous() for k, v in fx.items()} +save_file(fx, os.path.join(OUT_DIR, "gemma4_text.safetensors"), + metadata={"purpose": "diffuseR Gemma 4 test fixture"}) +print(f"wrote {len(fx)} tensors to {OUT_DIR}/gemma4_text.safetensors") + +# Tiny checkpoint in the layout of the LTX-2.5 bundled encoder file: +# model.* text weights, the connectors' text projection, a vision tensor +# the text-only loader must skip, and gemma_config metadata +ckpt = {f"model.{k}": v.detach().clone().contiguous() + for k, v in model.state_dict().items()} +ckpt["text_embedding_projection.video_aggregate_embed.weight"] = torch.randn(8, 32 * 7) +ckpt["text_embedding_projection.video_aggregate_embed.bias"] = torch.randn(8) +ckpt["vision_model.patch_dense.weight"] = torch.randn(4, 4) +gemma_config = {"architectures": ["Gemma4UnifiedForConditionalGeneration"], + "gemma_version": "gemma4-12b-ltx-v1", + "text_config": text_config} + +# --- Tokenizer ------------------------------------------------------------------ + +# Pieces: specials at the Gemma ids, the 256 byte-fallback tokens, single +# characters, and merged pieces reachable through the merge list +specials = ["", "", "", "", ""] +chars = list("abcdehilmnorstw.") + ["▁"] +merges = [("▁", "t"), ("h", "e"), ("▁t", "he"), ("l", "l"), + ("e", "ll"), ("o", "w"), ("▁", "c"), ("▁c", "a"), + ("a", "t"), ("▁ca", "t"), ("▁", "s"), ("i", "t"), + ("▁s", "it"), ("▁", "o"), ("n", "▁")] +vocab = {} +for s in specials: + vocab[s] = len(vocab) +for b in range(256): + vocab[f"<0x{b:02X}>"] = len(vocab) +for c in chars: + vocab.setdefault(c, len(vocab)) +for a, b in merges: + vocab.setdefault(a + b, len(vocab)) + +tok_json = { + "version": "1.0", + "truncation": None, + "padding": None, + "added_tokens": [{"id": vocab[s], "content": s, "single_word": False, + "lstrip": False, "rstrip": False, "normalized": False, + "special": True} for s in specials], + "normalizer": {"type": "Replace", "pattern": {"String": " "}, + "content": "▁"}, + "pre_tokenizer": {"type": "Split", "pattern": {"String": " "}, + "behavior": "MergedWithPrevious", "invert": False}, + "post_processor": {"type": "TemplateProcessing", + "single": [{"Sequence": {"id": "A", "type_id": 0}}], + "pair": [{"Sequence": {"id": "A", "type_id": 0}}, + {"Sequence": {"id": "B", "type_id": 1}}], + "special_tokens": {}}, + "decoder": None, + "model": {"type": "BPE", "dropout": None, "unk_token": "", + "continuing_subword_prefix": None, "end_of_word_suffix": None, + "fuse_unk": True, "byte_fallback": True, "ignore_merges": False, + "vocab": vocab, "merges": [[a, b] for a, b in merges]}, +} +tok_path = os.path.join(OUT_DIR, "gemma4_tiny_tokenizer.json") +with open(tok_path, "w", encoding="utf-8") as f: + json.dump(tok_json, f, ensure_ascii=False) + +# The LTX file carries its tokenizer.json as a uint8 tensor +with open(tok_path, "rb") as f: + tok_bytes = f.read() +ckpt["tokenizer_json"] = torch.tensor(list(tok_bytes), dtype=torch.uint8) +save_file(ckpt, os.path.join(OUT_DIR, "gemma4_tiny_ckpt.safetensors"), + metadata={"gemma_config": json.dumps(gemma_config), "format": "pt"}) +print(f"wrote {len(ckpt)} tensors to {OUT_DIR}/gemma4_tiny_ckpt.safetensors") + +tok = Tokenizer.from_file(tok_path) +prompts = ["the cat sat", "hello cat", " the cat ", "tell", "caté ok", + "thecat", "sit on the mat."] +cases = [{"text": p, "ids": tok.encode(p, add_special_tokens=True).ids} + for p in prompts] +with open(os.path.join(OUT_DIR, "gemma4_tokenizer_cases.json"), "w", + encoding="utf-8") as f: + json.dump(cases, f, ensure_ascii=False) +print(f"wrote {len(cases)} tokenizer cases") + +# Optional: ids from the real LTX-2.5 tokenizer, for an at-home R test +real = os.environ.get("GEMMA4_TOKENIZER_JSON") +if real and os.path.exists(real): + rtok = Tokenizer.from_file(real) + rprompts = [ + "A golden retriever runs across a sunlit meadow, slow motion.", + "Close-up: rain drips off a neon sign that reads \"OPEN 24/7\" — cinematic.", + "一只猫在窗台上睡觉。", + " two spaces and a tab\there ", + "Emoji test \U0001F680\U0001F319 and café crème brûlée", + ] + rcases = [{"text": p, "ids": rtok.encode(p, add_special_tokens=True).ids} + for p in rprompts] + with open(os.path.join(OUT_DIR, "gemma4_real_tokenizer_cases.json"), "w", + encoding="utf-8") as f: + json.dump(rcases, f, ensure_ascii=False) + print(f"wrote {len(rcases)} real-tokenizer cases")