LTX chain hardening: pinned-memory allocation, onload idempotency, connector_embeds, pinned Gemma3

DESCRIPTION

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

NEWS.md

@@ -1,3 +1,26 @@
+# diffuseR 0.1.0.15 (development)
+
+* Pinned staging now allocates page-locked host memory via
+  `torch_empty_strided(pin_memory = TRUE)` instead of the deprecated
+  `Tensor$pin_memory(device)` overload, silencing the two libtorch
+  deprecation warnings per tensor (thousands of lines per pipeline
+  load). Validated at 25.1 GB/s H2D (the card's DMA ceiling) with a
+  clean stderr.
+* `txt2vid_ltx2()` gains `connector_embeds=`: precomputed
+  text-connector outputs (the prompt is constant across a chained
+  track) skip the per-call connectors phase. Benchmarked at 960x960
+  this cuts the denoise-phase peak by ~2.5 GiB — the per-call phase
+  moved the raw Gemma3 hidden-state stack to the GPU, which is what
+  OOMed every resident-transformer chain attempt.
+* The `resident=` onload check now probes the staging pair's live
+  tensor instead of `module$parameters` (which the NF4 transformer
+  doesn't expose), so a resident component is never silently
+  re-onloaded over itself.
+* The Gemma3 loaders gain `pin=` (default: the `diffuseR.pin_staging`
+  option): a CPU-resident encoder is page-locked once and
+  `encode_with_gemma3()` swaps it to the GPU per encode (~0.3 s on,
+  free off) instead of holding VRAM or reloading (~15 s).
+
 # diffuseR 0.1.0.14 (development)
 
 * Latent-space chaining seams for chunked video continuation:

R/fp8_ltx23.R

@@ -54,7 +54,7 @@ ltx23_fp8_linear <- torch::nn_module(
                                      set_fp8_weight = function(weight, scale, pin = FALSE) {
     weight <- weight$to(device = "cpu")
     if (pin && torch::cuda_is_available()) {
-        weight <- weight$pin_memory(device = torch::torch_device("cuda"))
+        weight <- .ltx23_pin_host(weight)
     }
     self$weight_fp8 <- weight
     self$weight_scale <- scale$to(device = "cpu",

R/gemma3_text_encoder.R

@@ -581,19 +581,26 @@ gemma3_config_ltx2 <- function() {
 #' @param model_path Character. Path to directory containing model files.
 #' @param device Character. Device to load model to.
 #' @param dtype Character. Data type ("float32", "float16", "bfloat16").
+#' @param pin Logical. When loading to the CPU, page-lock the weights
+#'   so \code{\link{encode_with_gemma3}} can swap the model to the GPU
+#'   at DMA speed per encode (~0.3 s on, free off) instead of holding
+#'   VRAM or reloading. Default follows
+#'   \code{options(diffuseR.pin_staging)}.
 #' @param verbose Logical. Print loading progress.
 #' @return Initialized gemma3_text_model with loaded weights.
 #' @export
 load_gemma3_text_encoder <- function(model_path, device = "cpu",
-                                     dtype = "float16", verbose = TRUE) {
+                                     dtype = "float16",
+                                     pin = getOption("diffuseR.pin_staging", TRUE),
+                                     verbose = TRUE) {
     # An NF4 artifact directory (gemma3_quantize_nf4) loads through the
     # quantized path; bf16 compute unless the caller says otherwise
     if (file.exists(file.path(model_path, "manifest.json"))) {
         if (identical(dtype, "float16")) {
             dtype <- "bfloat16"
         }
         return(load_gemma3_nf4(model_path, device = device, dtype = dtype,
-                               verbose = verbose))
+                               pin = pin, verbose = verbose))
     }
     # Load config
     config_path <- file.path(model_path, "config.json")
@@ -679,6 +686,16 @@ load_gemma3_text_encoder <- function(model_path, device = "cpu",
 
     model$to(device = device)
 
+    if (pin && device == "cpu") {
+        st <- .ltx23_pin_component(model)
+        if (!is.null(st)) {
+            attr(model, "staging") <- st
+            if (verbose) {
+                message("Gemma3 weights pinned for staged transfer")
+            }
+        }
+    }
+
     if (verbose) {
         message("Gemma3 text encoder loaded on device: ", device)
     }
@@ -838,6 +855,17 @@ encode_with_gemma3 <- function(prompts, model = NULL, tokenizer = NULL,
         stop("Both model and tokenizer are required")
     }
 
+    # A pinned CPU-resident model (see the loaders' pin argument) swaps
+    # to the compute device for the encode and back for free afterwards
+    staging <- attr(model, "staging")
+    if (!is.null(staging) && device != "cpu") {
+        cur <- tryCatch(staging[[1]]$live$device$type, error = function(e) NULL)
+        if (!identical(cur, device)) {
+            .ltx23_staged_onload(staging, device)
+            on.exit(.ltx23_staged_offload(staging), add = TRUE)
+        }
+    }
+
     # Ensure prompts is a list
     if (is.character(prompts)) {
         prompts <- as.list(prompts)

R/quantize_gemma3.R

@@ -143,14 +143,21 @@ gemma3_quantize_nf4 <- function(model_path, output_dir = NULL,
 #'   \code{\link{gemma3_quantize_nf4}}.
 #' @param device "cuda" or "cpu".
 #' @param dtype Compute dtype ("bfloat16" default).
+#' @param pin Logical. When loading to the CPU, page-lock the weights
+#'   so \code{\link{encode_with_gemma3}} can swap the model to the GPU
+#'   at DMA speed per encode and back for free (see
+#'   \code{\link{staging_ltx23}}). Default follows
+#'   \code{options(diffuseR.pin_staging)}.
 #' @param verbose Logical.
 #'
 #' @return A \code{gemma3_text_model} ready for
 #'   \code{\link{encode_with_gemma3}}.
 #'
 #' @export
 load_gemma3_nf4 <- function(artifact_dir, device = "cuda",
-                            dtype = "bfloat16", verbose = TRUE) {
+                            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"))
     if (!identical(manifest$model, "gemma3")) {
@@ -208,6 +215,15 @@ load_gemma3_nf4 <- function(artifact_dir, device = "cuda",
 
     model$to(device = device)
     model$eval()
+    if (pin && device == "cpu") {
+        st <- .ltx23_pin_component(model)
+        if (!is.null(st)) {
+            attr(model, "staging") <- st
+            if (verbose) {
+                message("Gemma3 weights pinned for staged transfer")
+            }
+        }
+    }
     if (verbose) {
         message("Gemma3 NF4 encoder ready on ", device)
     }

R/staging_ltx23.R

@@ -33,20 +33,41 @@ NULL
 #'   if pinning is unavailable (no CUDA, or page-locking failed).
 #'
 #' @keywords internal
+# Allocate pinned host memory without Tensor$pin_memory(): this torch
+# build's binding requires the deprecated device argument (omitting it
+# errors with "Expected a torch_device"; passing it prints two libtorch
+# deprecation warnings per tensor - thousands of lines per pipeline
+# load). torch_empty_strided is the one creation op that exposes
+# pin_memory, so pin by allocating and copying; fall back to the noisy
+# path on builds where that fails.
+.ltx23_pin_host <- function(p) {
+    tryCatch({
+        sz <- as.integer(p$shape)
+        st <- if (length(sz)) {
+            as.integer(rev(cumprod(c(1, rev(sz)[-length(sz)]))))
+        } else {
+            integer(0)
+        }
+        buf <- torch::torch_empty_strided(sz, st, dtype = p$dtype,
+            pin_memory = TRUE)
+        torch::with_no_grad(buf$copy_(p))
+        buf
+    }, error = function(e) {
+        suppressWarnings(p$pin_memory(device = torch::torch_device("cuda")))
+    })
+}
+
 .ltx23_pin_component <- function(module) {
     if (!torch::cuda_is_available()) {
         return(NULL)
     }
-    cuda_dev <- torch::torch_device("cuda")
     tryCatch({
         tensors <- c(module$parameters, module$buffers)
-        suppressWarnings(lapply(tensors, function(p) {
-            # The device argument is deprecated upstream but this
-            # torch build requires it
-            pinned <- p$pin_memory(device = cuda_dev)
+        lapply(tensors, function(p) {
+            pinned <- .ltx23_pin_host(p)
             p$set_data(pinned)
             list(live = p, pinned = pinned)
-        }))
+        })
     }, error = function(e) NULL)
 }
 

R/txt2vid_ltx23.R

@@ -391,6 +391,14 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda",
 #' @param prompt_embeds Optional precomputed list with
 #'   \code{prompt_embeds} (raw stacked Gemma3 states) and
 #'   \code{prompt_attention_mask}; bypasses the text encoder.
+#' @param connector_embeds Optional precomputed text-connector outputs:
+#'   a list with \code{video_text_embedding},
+#'   \code{audio_text_embedding}, and \code{attention_mask} (the result
+#'   of \code{pipeline$connectors} on the Gemma3 states). The prompt is
+#'   constant across a chained track, so compute this once and pass it
+#'   to every chunk: it skips the per-call connectors phase, whose
+#'   GPU-side handling of the raw hidden-state stack does not fit next
+#'   to a resident transformer.
 #' @param width,height Integers. Output resolution (multiples of 32).
 #' @param num_frames Integer. 8k + 1 frames (e.g. 121).
 #' @param frame_rate Numeric. Frames per second.
@@ -471,8 +479,9 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda",
 #' @export
 txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
                          tokenizer = NULL, prompt_embeds = NULL,
-                         width = 768L, height = 512L, num_frames = 121L,
-                         frame_rate = 24, sigmas = ltx23_distilled_sigmas(),
+                         connector_embeds = NULL, width = 768L,
+                         height = 512L, num_frames = 121L, frame_rate = 24,
+                         sigmas = ltx23_distilled_sigmas(),
                          guidance_scale = 1, seed = NULL, device = "cuda",
                          dtype = "bfloat16", filename = NULL,
                          max_sequence_length = 1024L, decode_video = TRUE,
@@ -541,7 +550,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
     f32 <- torch::torch_float32()
 
     # --- Phase 1: text encoding -------------------------------------------------
-    if (is.null(prompt_embeds)) {
+    if (is.null(prompt_embeds) && is.null(connector_embeds)) {
         if (is.null(text_encoder) || is.null(tokenizer)) {
             stop("Supply text_encoder + tokenizer, or precomputed prompt_embeds.")
         }
@@ -573,9 +582,25 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
         }
         if (phase_offload) {
             # Idempotent: a resident component is already in place on
-            # the second and later calls of a chained run
-            cur <- tryCatch(module$parameters[[1]]$device$type,
-                            error = function(e) NULL)
+            # the second and later calls of a chained run. Probe the
+            # staging pair's live tensor when staging exists - custom
+            # module classes (the NF4 transformer) may not expose
+            # $parameters, and a failed probe must not degrade into a
+            # re-onload: re-transferring resident weights over
+            # themselves fragments the allocator pool until the next
+            # large allocation OOMs.
+            if (is.character(what)) {
+                st_probe <- staging[[what]]
+            } else {
+                st_probe <- NULL
+            }
+            cur <- tryCatch({
+                if (!is.null(st_probe)) {
+                    st_probe[[1]]$live$device$type
+                } else {
+                    module$parameters[[1]]$device$type
+                }
+            }, error = function(e) NULL)
             if (identical(cur, target_type)) {
                 return(module)
             }
@@ -622,19 +647,33 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL,
         invisible(module)
     }
 
-    onload("connectors")
-    connectors_device <- pipeline$connectors$video_text_proj_in$weight$device
-    torch::with_no_grad({
-        conn <- pipeline$connectors(
-                                    prompt_embeds$prompt_embeds$to(device = connectors_device, dtype = compute_dtype),
-                                    prompt_embeds$prompt_attention_mask$to(device = connectors_device)
-        )
-        video_text_embeds <- conn$video_text_embedding$to(device = device)
-        audio_text_embeds <- conn$audio_text_embedding$to(device = device)
-        text_mask <- conn$attention_mask$to(device = device)
-    })
-    rm(conn)
-    offload("connectors")
+    if (!is.null(connector_embeds)) {
+        # Precomputed connector outputs (constant per prompt): skip the
+        # per-call connectors phase entirely. Chained generation with a
+        # resident transformer needs this - the raw Gemma3 hidden-state
+        # stack moving to the GPU next to 11 GB of parked weights is
+        # what OOMed every resident chunk (geometry-independent 14.97
+        # GiB signature).
+        video_text_embeds <- connector_embeds$video_text_embedding$
+        to(device = device, dtype = compute_dtype)
+        audio_text_embeds <- connector_embeds$audio_text_embedding$
+        to(device = device, dtype = compute_dtype)
+        text_mask <- connector_embeds$attention_mask$to(device = device)
+    } else {
+        onload("connectors")
+        connectors_device <- pipeline$connectors$video_text_proj_in$weight$device
+        torch::with_no_grad({
+            conn <- pipeline$connectors(
+                                        prompt_embeds$prompt_embeds$to(device = connectors_device, dtype = compute_dtype),
+                                        prompt_embeds$prompt_attention_mask$to(device = connectors_device)
+            )
+            video_text_embeds <- conn$video_text_embedding$to(device = device)
+            audio_text_embeds <- conn$audio_text_embedding$to(device = device)
+            text_mask <- conn$attention_mask$to(device = device)
+        })
+        rm(conn)
+        offload("connectors")
+    }
     gc(verbose = FALSE)
 
     # --- Phase 2: latent preparation ---------------------------------------------

inst/tinytest/test_gemma3_loader.R

@@ -73,3 +73,10 @@ expect_error(
 )
 
 unlink(dir, recursive = TRUE)
+
+# --- pinned staging surface ----------------------------------------------------------
+
+# Both loaders take pin (default: the pin_staging option). Actual
+# page-locking needs CUDA; validated on GPU outside R CMD check.
+expect_true("pin" %in% names(formals(load_gemma3_text_encoder)))
+expect_true("pin" %in% names(formals(load_gemma3_nf4)))

inst/tinytest/test_txt2vid_ltx23.R

@@ -288,6 +288,36 @@ expect_error(
 expect_error(ltx23_tail_latents(res_cont$latents), pattern = "latent_shape")
 expect_error(ltx23_tail_latents(res_cont, k = 5L), pattern = "k must be")
 
+# --- connector_embeds bypass ---------------------------------------------------------
+
+# Precomputed connector outputs reproduce the internal connectors path
+# exactly, and satisfy the text requirement without prompt_embeds
+conn_out <- torch::with_no_grad(connectors(
+  stub_embeds$prompt_embeds, stub_embeds$prompt_attention_mask))
+cemb <- list(video_text_embedding = conn_out$video_text_embedding,
+             audio_text_embedding = conn_out$audio_text_embedding,
+             attention_mask = conn_out$attention_mask)
+res_ce <- txt2vid_ltx2(
+  prompt = "tiny bypass",
+  pipeline = pipe,
+  connector_embeds = cemb,
+  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
+  seed = 7L, device = "cpu", dtype = "float32",
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+res_full <- txt2vid_ltx2(
+  prompt = "tiny bypass",
+  pipeline = pipe,
+  prompt_embeds = stub_embeds,
+  width = 64L, height = 64L, num_frames = 9L, frame_rate = 24,
+  seed = 7L, device = "cpu", dtype = "float32",
+  decode_audio = FALSE,
+  verbose = FALSE
+)
+expect_equal(dim(res_ce$video), c(9L, 64L, 64L, 3L))
+expect_true(max(abs(res_ce$video - res_full$video)) < 1e-6)
+
 
 # --- Audio-conditioned generation (lip-sync plumbing) --------------------------------
 

man/dot-ltx23_pin_component.Rd → man/dot-ltx23_pin_host.Rd

@@ -1,9 +1,9 @@
 % tinyrox says don't edit this manually, but it can't stop you!
-\name{.ltx23_pin_component}
-\alias{.ltx23_pin_component}
+\name{.ltx23_pin_host}
+\alias{.ltx23_pin_host}
 \title{Pin a component's tensors for fast phase transfer}
 \usage{
-.ltx23_pin_component(module)
+.ltx23_pin_host(p)
 }
 \arguments{
 \item{module}{An nn_module on the CPU.}

man/load_gemma3_nf4.Rd

@@ -4,7 +4,7 @@
 \title{Load a Gemma3 text encoder from an NF4 artifact}
 \usage{
 load_gemma3_nf4(artifact_dir, device = "cuda", dtype = "bfloat16",
-                verbose = TRUE)
+                pin = getOption("diffuseR.pin_staging", TRUE), verbose = TRUE)
 }
 \arguments{
 \item{artifact_dir}{Directory produced by
@@ -14,6 +14,12 @@ load_gemma3_nf4(artifact_dir, device = "cuda", dtype = "bfloat16",
 
 \item{dtype}{Compute dtype ("bfloat16" default).}
 
+\item{pin}{Logical. When loading to the CPU, page-lock the weights
+so \code{\link{encode_with_gemma3}} can swap the model to the GPU
+at DMA speed per encode and back for free (see
+\code{\link{staging_ltx23}}). Default follows
+\code{options(diffuseR.pin_staging)}.}
+
 \item{verbose}{Logical.}
 }
 \value{