diff --git a/DESCRIPTION b/DESCRIPTION index c047ea6..53f5639 100644 --- a/DESCRIPTION +++ b/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")), diff --git a/NEWS.md b/NEWS.md index 04a8ae9..1750bc5 100644 --- a/NEWS.md +++ b/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//`. 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 diff --git a/R/gemma3_text_encoder.R b/R/gemma3_text_encoder.R index 4ac4de5..0d977f3 100644 --- a/R/gemma3_text_encoder.R +++ b/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 diff --git a/R/resident.R b/R/resident.R index 7a61b38..0ace6d7 100644 --- a/R/resident.R +++ b/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 diff --git a/R/staging.R b/R/staging.R index b5b5cea..8705ad4 100644 --- a/R/staging.R +++ b/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) diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R index b52b915..b3ba0c6 100644 --- a/R/txt2img_flux2.R +++ b/R/txt2img_flux2.R @@ -11,13 +11,46 @@ #' @name txt2img_flux2 NULL +# The revision a hub lookup resolves against. +# +# NULL means hfhub's default, "main" -- a BRANCH, which it resolves through +# `refs/main` in the cache and, failing that, over the network. An exact 40-hex +# commit takes hfhub's fast path instead: straight to +# `snapshots//`, never consulting `refs/`. +# +# That difference is what lets a pipeline load from a cache holding ONLY the +# snapshot -- a read-only bind of one revision, with no `refs/` and no network, +# which is how a fleet node serves weights it did not bake. It is also the +# stronger guarantee generally: a branch moves, so a deployment pinned to +# yesterday's weights would silently start resolving today's. +# +# A branch name is REFUSED rather than passed through. Accepting one would hand +# hfhub a value it resolves the slow way, which is the behaviour this argument +# exists to avoid, and the caller would have no way to tell. +.flux2_rev <- function(revision) { + if (is.null(revision)) { + return(list()) + } + if (!is.character(revision) || length(revision) != 1L || is.na(revision) || + !grepl("^[0-9a-f]{40}$", revision)) { + stop("revision must be a single 40-character hex commit, not a branch ", + "name: a branch resolves through refs/ and defeats the point of ", + "pinning one", call. = FALSE) + } + list(revision = revision) +} + # Resolve a FLUX.2-klein support file from the HuggingFace cache -.flux2_cached <- function(file) { +.flux2_cached <- function(file, revision = NULL) { if (!requireNamespace("hfhub", quietly = TRUE)) { stop("The hfhub package is required to locate model files.") } + ## Outside the tryCatch below, which reports every error as a missing + ## download: a branch name is a caller mistake, not an absent file. + rev <- .flux2_rev(revision) tryCatch( - hfhub::hub_download(.flux2_repo, file, local_files_only = TRUE), + do.call(hfhub::hub_download, + c(list(.flux2_repo, file, local_files_only = TRUE), rev)), error = function(e) { stop("Missing ", file, " in the HuggingFace cache; ", "run download_flux2_klein() first.", call. = FALSE) @@ -47,6 +80,12 @@ NULL #' resolves via \code{options(diffuseR.pin_staging)} then the #' host-RAM-aware \code{\link{recommend}} decision. #' @param verbose Logical. +#' @param revision Optional exact 40-hex commit for the support files (VAE, +#' Qwen3 encoder, tokenizer) this pulls from the \code{black-forest-labs} +#' cache. With one they resolve straight out of +#' \code{snapshots//}, so the load works against a cache holding +#' only that snapshot -- no \code{refs/} entry and no network. The +#' transformer comes from \code{model_dir} and is unaffected. #' #' @return A \code{flux2_pipeline} list. #' @@ -55,7 +94,10 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4", "bf16"), text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, pin = NULL, - verbose = TRUE) { + verbose = TRUE, revision = NULL) { + ## Checked before any device work, where a configuration mistake is + ## cheapest to report. + .flux2_rev(revision) precision <- .flux_resolve_precision(match.arg(precision), file.path(tools::R_user_dir("diffuseR", "data"), "flux2-klein-4b-")) if (is.null(text_device)) { @@ -124,10 +166,11 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading FLUX.2 VAE decoder...") } - vae_config <- jsonlite::fromJSON(.flux2_cached("vae/config.json")) + vae_config <- jsonlite::fromJSON(.flux2_cached("vae/config.json", + revision)) pipe$vae_bn_eps <- vae_config$batch_norm_eps %||% 1e-4 pipe$decoder <- load_flux2_vae_decoder( - .flux2_cached("vae/diffusion_pytorch_model.safetensors"), + .flux2_cached("vae/diffusion_pytorch_model.safetensors", revision), latent_channels = as.integer(vae_config$latent_channels %||% 32L), verbose = verbose ) @@ -142,13 +185,14 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading Qwen3 text encoder...") } - te_dir <- dirname(.flux2_cached("text_encoder/config.json")) + te_dir <- dirname(.flux2_cached("text_encoder/config.json", revision)) pipe$text_encoder <- load_qwen3_text_encoder( te_dir, device = if (phase_offload) "cpu" else text_device, dtype = if (text_device == "cpu") "float32" else "bfloat16", verbose = verbose ) - pipe$tokenizer <- qwen_bpe_tokenizer(.flux2_cached("tokenizer/tokenizer.json")) + pipe$tokenizer <- qwen_bpe_tokenizer( + .flux2_cached("tokenizer/tokenizer.json", revision)) components <- c("transformer", "decoder") if (!identical(text_device, "cpu")) { diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 8501d7d..7767cc4 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -373,6 +373,32 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", structure(pipe, class = "ltx23_pipeline") } +# The device the prompt encode runs on, for the text_encoder txt2vid_ltx2 +# was handed. A PATH is loaded fresh onto `device` -- a caller who names a +# file is asking for it there. A PRELOADED object is CPU-resident by the +# resident loader's convention (load_gemma3_text_encoder(device = "cpu", +# pin = TRUE)), and only safe to encode on the compute device when it +# carries the pinned `staging` set that lets encode_with_gemma3() DMA it +# there and back per encode. Without that set, `device = "cuda"` would send +# the tokens to the card while the weights sat on the host -- a mismatch -- +# so a bare object degrades to CPU. +# +# The `else "cpu"` this replaces was blunter than that: it forced CPU for +# EVERY preloaded encoder, so the resident/gpuhost path (which loads with +# pin = TRUE precisely to stage) encoded every prompt on CPU (~24 s) with +# its own pinned staging sitting unused, instead of the ~7 s staged-GPU +# encode. Split out so the decision is assertable without a GPU or a real +# encode, the way `.resident_gen_args` is. +.ltx23_text_encode_device <- function(text_encoder, device) { + if (is.character(text_encoder)) { + return(device) + } + if (!is.null(attr(text_encoder, "staging")) && grepl("^cuda", device)) { + return(device) + } + "cpu" +} + #' Generate video (and audio) with LTX-2.3 #' #' Distilled text-to-video generation: encodes the prompt with Gemma3 + @@ -561,7 +587,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, prompt, model = text_encoder, tokenizer = tokenizer, max_sequence_length = max_sequence_length, - device = if (is.character(text_encoder)) device else "cpu", + device = .ltx23_text_encode_device(text_encoder, device), verbose = verbose ) } @@ -582,37 +608,30 @@ 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. Probe the - # staging pair's live tensor when staging exists - custom - # module classes (the NF4 transformer) may not expose + # the second and later calls of a chained run, and must not + # be re-transferred over itself (that fragments the + # allocator pool until the next large allocation OOMs). + # With staging, `.staged_onload` decides PER PAIR: a + # resident component is a no-op, and one whose earlier + # onload stopped partway is completed instead of being + # reported resident by its first tensor and left split. + # Without staging, probe the module -- 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) - } + # re-onload. if (is.character(what)) { st <- staging[[what]] } else { st <- NULL } - if (is.null(st)) { - module$to(device = device) - } else { + if (!is.null(st)) { .staged_onload(st, device) + return(module) + } + cur <- tryCatch(module$parameters[[1]]$device$type, + error = function(e) NULL) + if (!identical(cur, target_type)) { + module$to(device = device) } } module diff --git a/inst/tinytest/test_resident.R b/inst/tinytest/test_resident.R index c5024ca..ec74cd6 100644 --- a/inst/tinytest/test_resident.R +++ b/inst/tinytest/test_resident.R @@ -166,6 +166,27 @@ expect_error(resident_activate(u), pattern = "unloaded") # Status still works on an unloaded handle. expect_equal(resident_status(u)$state, "unloaded") +# Unload drops the LTX text encoder too. It lives outside `staging` on +# purpose (resident_load), so a handle that only cleared `staging` kept +# the encoder's pinned buffers while reporting pinned_bytes = 0. +# A fake pair records the offload it receives. +te_log <- new.env() +te_log$calls <- list() +te_pair <- list(live = list(set_data = function(x) { + te_log$calls[[length(te_log$calls) + 1L]] <- x + }), + pinned = "pinned-copy") +u2 <- mk("inactive") +u2$text_encoder <- structure(list(), staging = list(te_pair)) +u2$tokenizer <- list(vocab = 3L) +resident_unload(u2) +expect_equal(u2$state, "unloaded") +expect_null(u2$text_encoder) +expect_null(u2$tokenizer) +# The encoder's pinned staging was offloaded (pointer swap back to the +# pinned copy) before being dropped. +expect_equal(te_log$calls, list("pinned-copy")) + # --- CUDA round trip -------------------------------------------------------------- if (have_torch && at_home() && torch::cuda_is_available()) { diff --git a/inst/tinytest/test_staged_on.R b/inst/tinytest/test_staged_on.R new file mode 100644 index 0000000..710a5cf --- /dev/null +++ b/inst/tinytest/test_staged_on.R @@ -0,0 +1,95 @@ +# .staged_on / .staged_onload (R/staging.R): the checks a phase makes +# before moving a pinned component. Pure fakes, no torch, no GPU: a +# "pair" is anything with $live$device, $live$set_data and $pinned$to, +# which is all the helpers touch. +# +# The case that matters is the PARTIAL one. An onload that dies partway +# leaves the first pairs on the card and the rest on the host. Probing +# the first pair alone called that resident, so every later phase +# skipped the onload and failed on a device mismatch until the process +# ended (the gpuhost's ltx-2.3 entry, 2026-09-10). + +library(diffuseR) +staged_on <- diffuseR:::.staged_on +staged_onload <- diffuseR:::.staged_onload +device_spec <- diffuseR:::.device_spec + +# A fake pair whose live tensor sits on `type` (and card `index`), +# recording every set_data it receives. +fake_pair <- function(type, index = if (type == "cpu") NULL else 0L) { + log <- new.env() + log$calls <- list() + live <- list(device = list(type = type, index = index), + set_data = function(x) { + log$calls[[length(log$calls) + 1L]] <- x + }) + pinned <- list(to = function(device, non_blocking = FALSE) { + paste0("copy->", device) + }) + list(live = live, pinned = pinned, log = log) +} +calls <- function(p) p$log$calls + +# Device spec: strings with and without an index, and a torch_device. +expect_equal(device_spec("cuda"), list(type = "cuda", index = NA_integer_)) +expect_equal(device_spec("cuda:1"), list(type = "cuda", index = 1L)) +expect_equal(device_spec("cpu"), list(type = "cpu", index = NA_integer_)) +expect_equal(device_spec(structure(list(type = "cuda", index = 1), + class = "torch_device")), + list(type = "cuda", index = 1L)) +expect_equal(device_spec(structure(list(type = "cuda", index = NULL), + class = "torch_device")), + list(type = "cuda", index = NA_integer_)) + +# All on the card: resident. +st <- list(fake_pair("cuda"), fake_pair("cuda"), fake_pair("cuda")) +expect_true(staged_on(st, "cuda")) +expect_true(staged_on(st, "cuda:0")) +expect_false(staged_on(st, "cpu")) + +# All on the host: not resident. +st <- list(fake_pair("cpu"), fake_pair("cpu")) +expect_false(staged_on(st, "cuda")) +expect_true(staged_on(st, "cpu")) + +# PARTIAL: first pair on the card, second on the host. The first-pair +# probe said "resident"; the whole-set check must not. +st <- list(fake_pair("cuda"), fake_pair("cpu")) +expect_false(staged_on(st, "cuda")) + +# THE WRONG CARD is not this card. A request naming an index accepts only +# that index; a request without one accepts any card. +st <- list(fake_pair("cuda", 0L), fake_pair("cuda", 1L)) +expect_false(staged_on(st, "cuda:0")) +expect_false(staged_on(st, "cuda:1")) +expect_true(staged_on(st, "cuda")) +expect_true(staged_on(list(fake_pair("cuda", 1L)), "cuda:1")) + +# An unreadable pair is not resident. +st <- list(fake_pair("cuda"), list(live = NULL, pinned = NULL)) +expect_false(staged_on(st, "cuda")) + +# Empty staging holds nothing to move. +expect_true(staged_on(list(), "cuda")) + +# Onload is idempotent per pair: resident pairs are untouched, the rest +# are copied. A partial onload is completed, not restarted. +st <- list(fake_pair("cuda"), fake_pair("cpu"), + fake_pair("cuda"), fake_pair("cpu")) +staged_onload(st, "cuda") +expect_equal(length(calls(st[[1]])), 0L) +expect_equal(calls(st[[2]]), list("copy->cuda")) +expect_equal(length(calls(st[[3]])), 0L) +expect_equal(calls(st[[4]]), list("copy->cuda")) + +# A fully resident set is a no-op, whichever spelling names the card. +st <- list(fake_pair("cuda"), fake_pair("cuda")) +staged_onload(st, "cuda:0") +expect_equal(length(calls(st[[1]])), 0L) +expect_equal(length(calls(st[[2]])), 0L) + +# A pair on another card IS moved when a particular card was asked for. +st <- list(fake_pair("cuda", 0L), fake_pair("cuda", 1L)) +staged_onload(st, "cuda:1") +expect_equal(calls(st[[1]]), list("copy->cuda:1")) +expect_equal(length(calls(st[[2]])), 0L) diff --git a/inst/tinytest/test_staging.R b/inst/tinytest/test_staging.R index 0490091..0c64e4f 100644 --- a/inst/tinytest/test_staging.R +++ b/inst/tinytest/test_staging.R @@ -43,3 +43,20 @@ diffuseR:::.staged_onload(st, "cuda") torch::with_no_grad(out_gpu2 <- m(x$to(device = "cuda"))$cpu()) expect_true(as.numeric((out_gpu2 - out_gpu)$abs()$max()) == 0) diffuseR:::.staged_offload(st) + +# PARTIAL onload, the state a failed transfer leaves behind: the first +# pair on the card, the rest on the host. The whole-set check must call +# it not resident, and the next onload must complete it -- moving only +# what is missing -- rather than skip it on the strength of pair 1. +st[[1]]$live$set_data(st[[1]]$pinned$to(device = "cuda")) +expect_equal(st[[1]]$live$device$type, "cuda") +expect_equal(st[[length(st)]]$live$device$type, "cpu") +expect_false(diffuseR:::.staged_on(st, "cuda")) +diffuseR:::.staged_onload(st, "cuda") +expect_true(diffuseR:::.staged_on(st, "cuda")) +torch::with_no_grad(out_gpu3 <- m(x$to(device = "cuda"))$cpu()) +expect_true(as.numeric((out_gpu3 - out_gpu)$abs()$max()) == 0) +diffuseR:::.staged_offload(st) +expect_true(diffuseR:::.staged_on(st, "cpu")) +torch::with_no_grad(out_back2 <- m(x)) +expect_true(as.numeric((out_back2 - ref)$abs()$max()) == 0) diff --git a/inst/tinytest/test_text_encode_device.R b/inst/tinytest/test_text_encode_device.R new file mode 100644 index 0000000..44382f7 --- /dev/null +++ b/inst/tinytest/test_text_encode_device.R @@ -0,0 +1,31 @@ +# .ltx23_text_encode_device: which device the prompt encode runs on for +# the text_encoder txt2vid_ltx2 was handed. Pure decision, no torch, no +# GPU, no real encode -- the point is to pin the rule that a resident +# (preloaded, pinned) encoder stages to the card instead of encoding on +# CPU, which is what the gpuhost/resident path silently did before. + +library(diffuseR) +f <- diffuseR:::.ltx23_text_encode_device + +# A PATH loads fresh onto whatever device was asked for. +expect_equal(f("/models/gemma3-nf4", "cuda"), "cuda") +expect_equal(f("/models/gemma3-nf4", "cuda:0"), "cuda:0") +expect_equal(f("/models/gemma3-nf4", "cpu"), "cpu") + +# A PRELOADED object is only encoded on the card when it carries the +# pinned staging set (the resident loader's pin = TRUE). This is the case +# that regressed: the resident/gpuhost encoder HAS staging and a cuda +# request, so it must stage to the card, not fall to CPU. +staged <- structure(list(), staging = list(TRUE)) +expect_equal(f(staged, "cuda"), "cuda") +expect_equal(f(staged, "cuda:0"), "cuda:0") +# An explicit cpu request is still honoured even with staging present. +expect_equal(f(staged, "cpu"), "cpu") + +# A preloaded object with NO staging cannot run on the card -- its weights +# sit on the host and a cuda request would be a device mismatch -- so it +# degrades to CPU. This is the safety the blunt `else "cpu"` was protecting, +# preserved for exactly this case. +bare <- structure(list()) +expect_equal(f(bare, "cuda"), "cpu") +expect_equal(f(bare, "cpu"), "cpu") diff --git a/man/dot-staged_on.Rd b/man/dot-staged_on.Rd new file mode 100644 index 0000000..90171e6 --- /dev/null +++ b/man/dot-staged_on.Rd @@ -0,0 +1,34 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.staged_on} +\alias{.staged_on} +\title{Is every pinned tensor of a component on this device?} +\usage{ +.staged_on(staging, device) +} +\arguments{ +\item{staging}{A component's staging set: the list of +\code{list(live, pinned)} pairs \code{.pin_component} returned.} + +\item{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.} +} +\value{ +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. +} +\description{ +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. +} +\keyword{internal} diff --git a/man/dot-staged_onload.Rd b/man/dot-staged_onload.Rd index b9bb780..ace3459 100644 --- a/man/dot-staged_onload.Rd +++ b/man/dot-staged_onload.Rd @@ -8,5 +8,13 @@ \description{ Non-blocking copies from pinned memory share the default stream, so later kernels are ordered after them; no explicit sync needed. +} +\details{ +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. + } \keyword{internal} diff --git a/man/flux2_load_pipeline.Rd b/man/flux2_load_pipeline.Rd index 0ea182a..809e978 100644 --- a/man/flux2_load_pipeline.Rd +++ b/man/flux2_load_pipeline.Rd @@ -11,7 +11,8 @@ flux2_load_pipeline( attn_chunk = NULL, phase_offload = TRUE, pin = NULL, - verbose = TRUE + verbose = TRUE, + revision = NULL ) } \arguments{ @@ -37,6 +38,13 @@ resolves via \code{options(diffuseR.pin_staging)} then the host-RAM-aware \code{\link{recommend}} decision.} \item{verbose}{Logical.} + +\item{revision}{Optional exact 40-hex commit for the support files (VAE, +Qwen3 encoder, tokenizer) this pulls from the \code{black-forest-labs} +cache. With one they resolve straight out of +\code{snapshots//}, so the load works against a cache holding +only that snapshot -- no \code{refs/} entry and no network. The +transformer comes from \code{model_dir} and is unaffected.} } \value{ A \code{flux2_pipeline} list. diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 8f0bcd1..808e95c 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -7,7 +7,9 @@ resident_load( model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), device = "cuda", ..., - verbose = TRUE + verbose = TRUE, + text_encoder = NULL, + tokenizer = NULL ) } \arguments{ @@ -23,6 +25,15 @@ resident_load( \code{\link{download_sdxl}} cache).} \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}} +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.} } \value{ A \code{diffuseR_resident} handle (an environment). Inspect it