LTX-2.5 distilled text-to-video

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")),

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)

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

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
             )
         }
     }

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"
+    )
+}

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