Pin FLUX.2's revision and fix the partial-onload wedge in resident LTX

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.2.2.7
+Version: 0.2.2.10
 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,72 @@
+# diffuseR 0.2.2.10
+
+* **A pinned component whose onload failed partway no longer stays
+  wedged.** The Gemma3 encoder's staged encode, and the LTX pipeline's
+  per-phase onload, decided "already on the card" by probing the FIRST
+  staging pair. An onload that dies partway -- device memory runs out
+  with most of the encoder copied -- leaves exactly that pair on the card
+  and the rest on the host, so every later call skipped the onload and
+  failed on the first matrix multiply with "mat2 is on cpu", on every
+  request, until the process ended. That is what took the gpuhost's
+  ltx-2.3 entry down for USA 20260912 on 2026-09-10. Three changes:
+  `.staged_on()` asks every pair, not the first; `.staged_onload()` is
+  idempotent per pair, so a resident component is a no-op (no re-transfer
+  over itself) and a partial one is completed; and `encode_with_gemma3()`
+  arms its offload BEFORE the onload, so a failed transfer is undone on
+  the way out and the next encode starts from a clean host copy. Covered
+  by `test_staged_on.R` (pure fakes, no GPU) and a partial round trip in
+  `test_staging.R` (CUDA).
+
+* **Staging compares the card, not just the device type.** A request
+  for `"cuda:1"` no longer counts a tensor on `cuda:0` as resident, so a
+  multi-GPU caller asking for a particular card gets its weights moved
+  there instead of a skipped transfer and a device mismatch. A request
+  for bare `"cuda"` still accepts any card, as before.
+
+* **`resident_unload()` drops the LTX text encoder.** The encoder a
+  resident LTX handle loads (0.2.2.8) sits outside `staging` by design,
+  and unload never released it: an unloaded handle kept the encoder's
+  pinned buffers while reporting `pinned_bytes = 0`. Both review findings
+  from the 2026-09-11 Codex pass.
+
+# diffuseR 0.2.2.9
+
+* **A resident LTX encoder now stages the prompt encode to the card
+  instead of running it on CPU.** `txt2vid_ltx2()` chose the encode device
+  with `if (is.character(text_encoder)) device else "cpu"` -- so a
+  PRELOADED encoder (the resident/gpuhost path, `resident_load("ltx",
+  text_encoder = ...)`) always encoded on CPU, even though the resident
+  loader page-locks it with `pin = TRUE` for exactly the staged transfer
+  `encode_with_gemma3()` supports. The pinned staging sat unused and every
+  prompt paid the ~24 s CPU encode instead of the ~7 s staged-GPU one; on
+  the gpuhost path that is once per chunk. The device decision is now
+  `.ltx23_text_encode_device()`: a path loads onto the asked-for device, a
+  preloaded encoder with a `staging` set and a cuda request stages to the
+  card, and a bare CPU-resident object still degrades to CPU (a cuda
+  request without staging would be a device mismatch). Pure and unit-tested
+  without a GPU (`test_text_encode_device.R`).
+
+# diffuseR 0.2.2.8
+
+* `flux2_load_pipeline()` takes a `revision`. Its VAE, Qwen3 encoder and
+  tokenizer come from the Hugging Face cache, and hfhub's default revision
+  is the branch `main` -- resolved through `refs/main` and, failing that,
+  over the network. A read-only bind of one snapshot carries neither, so
+  the load failed there; an exact 40-hex commit takes hfhub straight to
+  `snapshots/<revision>/<file>`. A branch name is refused rather than
+  passed through.
+
+* `resident_load("ltx", ...)` takes `text_encoder` and `tokenizer` paths.
+  `ltx23_load_pipeline()` does not load them and `txt2vid_ltx2()` takes
+  them per call, so a resident LTX handle could be activated and could not
+  generate -- and a serving caller passing paths re-read 7.6 GB of Gemma3
+  on every request. Given here they load once, pinned on the host, and
+  `resident_generate()` supplies them. They stay OUT of the handle's
+  staging on purpose: `resident_activate()` places everything in staging at
+  once, and the encoder does not fit beside the transformer. It rides to
+  the card for its own phase and back off, as the pipeline's components do.
+  `pinned_bytes` counts it.
+
 # diffuseR 0.2.2.7
 
 * `recommend()` diagnosed the wrong safetensors capability for bf16. The

R/gemma3_text_encoder.R

@@ -860,14 +860,24 @@ encode_with_gemma3 <- function(prompts, model = NULL, tokenizer = NULL,
     }
 
     # A pinned CPU-resident model (see the loaders' pin argument) swaps
-    # to the compute device for the encode and back for free afterwards
+    # to the compute device for the encode and back for free afterwards.
+    # A model already on the device IN FULL is left there, and left alone
+    # on exit: whoever put it there owns it.
+    #
+    # THE OFFLOAD IS ARMED BEFORE THE ONLOAD, not after it. An onload that
+    # fails partway -- device memory runs out with most of the encoder
+    # already copied -- must still be undone on the way out, or the
+    # encoder stays half on the card with nothing left to move it back,
+    # and every encode after it inherits the split. `.staged_on` asks
+    # every pair rather than the first one for the same reason: the first
+    # tensor of a half-copied encoder IS on the card, and a probe of it
+    # alone reported the encoder resident and skipped the onload on every
+    # request for the rest of the process (the gpuhost's ltx-2.3 entry,
+    # 2026-09-10: "mat2 is on cpu" from every encode after one failure).
     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)) {
-            .staged_onload(staging, device)
-            on.exit(.staged_offload(staging), add = TRUE)
-        }
+    if (!is.null(staging) && device != "cpu" && !.staged_on(staging, device)) {
+        on.exit(.staged_offload(staging), add = TRUE)
+        .staged_onload(staging, device)
     }
 
     # Ensure prompts is a list

R/resident.R

@@ -235,6 +235,14 @@
 #'   \code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the
 #'   \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}}
+#'   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
+#'   generate; the encoder rides to the card for its phase and back off, the
+#'   way the pipeline's own components do, so it is never resident beside the
+#'   transformer.
 #'
 #' @return A \code{diffuseR_resident} handle (an environment). Inspect it
 #'   with \code{\link{resident_status}}; the fields of interest are the
@@ -256,8 +264,19 @@
 #' @export
 resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx",
                                     "sdxl", "sd21"),
-                          device = "cuda", ..., verbose = TRUE) {
+                          device = "cuda", ..., verbose = TRUE,
+                          text_encoder = NULL, tokenizer = NULL) {
     model <- match.arg(model)
+    ## NAMED ARGUMENTS RATHER THAN `...`, because `...` goes to the family
+    ## loader and `ltx23_load_pipeline` has no `...` of its own -- an unknown
+    ## argument there is an error, not a pass-through. Refused for the other
+    ## families for the same reason: silently ignoring them would leave a
+    ## caller believing a text encoder had been loaded.
+    if (!identical(model, "ltx") &&
+        (!is.null(text_encoder) || !is.null(tokenizer))) {
+        stop("text_encoder/tokenizer apply to the ltx family only; ",
+             model, " loads its own", call. = FALSE)
+    }
     if (!torch::cuda_is_available()) {
         stop("resident_load() requires CUDA", call. = FALSE)
     }
@@ -308,6 +327,48 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx",
     res$gpu_components <- pipeline$gpu_components
     res$components <- names(.resident_components(pipeline))
     res$pinned_bytes <- .resident_pinned_bytes(staging)
+
+    ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING.
+    ##
+    ## `txt2vid_ltx2` takes `text_encoder` and `tokenizer` per call and
+    ## accepts a PATH, which it then loads -- so a serving caller that passed
+    ## paths would re-read 7.6 GB of Gemma3 on every request. Loading here
+    ## makes it once.
+    ##
+    ## Deliberately NOT added to `staging`: `resident_activate` puts
+    ## everything in staging on the card at once, and the encoder beside the
+    ## transformer does not fit. It carries its own staging attribute from
+    ## `pin = TRUE`, and `encode_with_gemma3` onloads it for the encode and
+    ## offloads on exit -- one GPU tenant per phase, the same discipline the
+    ## pipeline's own components follow.
+    if (identical(model, "ltx") && !is.null(text_encoder)) {
+        if (is.null(tokenizer)) {
+            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)
+        ## 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
+        ## of this number.
+        ## WRAPPED IN A LIST, AND THAT IS NOT COSMETIC. There are two
+        ## staging shapes in this package: a pipeline's is a list OF
+        ## COMPONENTS each holding a list of pairs, which is why
+        ## `.resident_pinned_bytes` loops twice; an encoder's
+        ## `attr(model, "staging")` is a FLAT list of pairs, which is why
+        ## `.staged_onload` loops once. Passing the flat one straight in
+        ## reads a pair's fields as pairs and dies on `pair$pinned$shape`
+        ## -- "object of type 'closure' is not subsettable", from inside a
+        ## worker, sixty seconds after the pin began.
+        te_staging <- attr(res$text_encoder, "staging")
+        if (!is.null(te_staging)) {
+            res$pinned_bytes <- res$pinned_bytes +
+                .resident_pinned_bytes(list(te_staging))
+        }
+    }
     res$state <- "inactive"
     res$last_error <- NULL
     res$loaded_at <- Sys.time()
@@ -762,6 +823,17 @@ resident_generate <- function(res, prompt, ...) {
             }
         }), want)
     }
+    ## LTX takes its text encoder per call and stores none, so a handle that
+    ## loaded one has to hand it over on every generate. An explicit argument
+    ## still wins -- this fills a gap rather than overriding a caller who
+    ## brought precomputed embeds or a different encoder.
+    if (identical(res$model, "ltx") && !is.null(res$text_encoder) &&
+        is.null(args$text_encoder) && is.null(args$prompt_embeds)) {
+        args$text_encoder <- res$text_encoder
+        if (is.null(args$tokenizer)) {
+            args$tokenizer <- res$tokenizer
+        }
+    }
     args
 }
 
@@ -831,8 +903,17 @@ resident_unload <- function(res) {
         for (nm in names(res$staging)) {
             .staged_offload(res$staging[[nm]])
         }
+        # The LTX text encoder is deliberately outside `staging` (see
+        # resident_load), but it is pinned host memory this handle owns
+        # and counted in pinned_bytes, so it goes the same way.
+        te_staging <- attr(res$text_encoder, "staging")
+        if (!is.null(te_staging)) {
+            .staged_offload(te_staging)
+        }
     }, error = function(e) NULL)
     res$pipeline <- NULL
+    res$text_encoder <- NULL
+    res$tokenizer <- NULL
     res$staging <- list()
     res$components <- character(0)
     res$pinned_bytes <- 0

R/staging.R

@@ -80,14 +80,92 @@ NULL
     }, error = function(e) NULL)
 }
 
+# What a device spec names, as type and index. The index is NA when the
+# spec leaves it open: "cuda" means whichever card is current, so a
+# target without an index accepts any card, and "cuda:1" accepts only
+# that one. A tensor's own device always carries a concrete index on the
+# card (torch reports 0 for "cuda"), so the comparison below is exact
+# whenever the caller asked for a particular card. Parsed from the string
+# rather than through torch_device() so `.staged_on` runs without torch,
+# which is how it is tested.
+.device_spec <- function(device) {
+    if (inherits(device, "torch_device")) {
+        return(list(type = device$type,
+                    index = as.integer(device$index %||% NA_integer_)))
+    }
+    s <- as.character(device)
+    index <- if (grepl(":", s, fixed = TRUE)) {
+        as.integer(sub("^[^:]*:", "", s))
+    } else {
+        NA_integer_
+    }
+    list(type = sub(":.*$", "", s), index = index)
+}
+
+# Is this tensor on the device the spec names? Type must match; the
+# index must match too when the spec has one.
+.on_device <- function(tensor, spec) {
+    d <- tryCatch(tensor$device, error = function(e) NULL)
+    if (is.null(d) || !identical(d$type, spec$type)) {
+        return(FALSE)
+    }
+    if (is.na(spec$index)) {
+        return(TRUE)
+    }
+    identical(as.integer(d$index %||% NA_integer_), spec$index)
+}
+
+#' Is every pinned tensor of a component on this device?
+#'
+#' The check a caller makes before skipping an onload. It asks EVERY
+#' pair, not the first one: a component is on the card when all of it
+#' is, and a probe of one tensor cannot tell a resident component from
+#' one whose onload failed partway. That partial state is real -- an
+#' onload that runs out of device memory leaves the pairs it copied on
+#' the card and the rest on the host -- and a first-pair probe reports
+#' it as "already resident", so every later phase skips the onload and
+#' dies on a device mismatch, on every call, until the process ends.
+#' That is how a gpuhost's LTX entry wedged for a whole show on
+#' 2026-09-10: one failed encoder onload, then "mat2 is on cpu" from
+#' every request after it.
+#'
+#' @param staging A component's staging set: the list of
+#'   \code{list(live, pinned)} pairs \code{.pin_component} returned.
+#' @param device The compute device, as a string (\code{"cuda"},
+#'   \code{"cuda:1"}) or a \code{torch_device}. A spec without an index
+#'   accepts any card; one with an index accepts only that card.
+#' @return TRUE when every pair's live tensor is on that device; FALSE
+#'   on any mismatch or unreadable pair. Vacuously TRUE for an empty
+#'   staging set, which holds nothing to move.
+#' @keywords internal
+.staged_on <- function(staging, device) {
+    spec <- .device_spec(device)
+    for (pair in staging) {
+        if (!.on_device(pair$live, spec)) {
+            return(FALSE)
+        }
+    }
+    TRUE
+}
+
 #' Move a pinned component onto the compute device
 #'
 #' Non-blocking copies from pinned memory share the default stream,
 #' so later kernels are ordered after them; no explicit sync needed.
 #'
+#' Idempotent PER PAIR: a tensor already on the device is left where it
+#' is, so a resident component costs nothing to onload again (no
+#' re-transfer of weights over themselves, which fragments the
+#' allocator pool) and a component whose earlier onload stopped partway
+#' is completed rather than restarted.
+#'
 #' @keywords internal
 .staged_onload <- function(staging, device) {
+    spec <- .device_spec(device)
     for (pair in staging) {
+        if (.on_device(pair$live, spec)) {
+            next
+        }
         pair$live$set_data(pair$pinned$to(device = device, non_blocking = TRUE))
     }
     invisible(NULL)