diff --git a/DESCRIPTION b/DESCRIPTION index 73c3b11..38abf16 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.1.1 +Version: 0.2.1.2 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 1dcbfba..abb78e0 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -181,6 +181,12 @@ export(qwen_bpe_tokenizer) export(qwen3_encoder) export(recommend) export(reshard_safetensors) +export(resident_activate) +export(resident_deactivate) +export(resident_generate) +export(resident_load) +export(resident_status) +export(resident_unload) export(save_frames) export(save_image) export(save_video) @@ -229,6 +235,7 @@ export(zimage_transformer) export(zimage_unpatchify) S3method(print,bpe_tokenizer) +S3method(print,diffuseR_resident) S3method(print,ltx23_checkpoint) S3method(print,qwen_tokenizer) S3method(print,unigram_tokenizer) diff --git a/NEWS.md b/NEWS.md index 0241eb0..2c8652f 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,29 @@ +# diffuseR 0.2.1.2 + +* Model residency: `resident_load()`, `resident_activate()`, + `resident_deactivate()`, `resident_generate()`, `resident_status()` + and `resident_unload()` keep a pipeline's weights page-locked on the + host and treat the GPU copy as disposable, so handing a small card + between models is a DMA transfer rather than a full reload. Same + contract as whisper and chatterbox, with no `gpu.ctl` dependency. + This sits above the per-generation phase offloading in the + `txt2img_*` functions: those swap one component at a time within a + render, residency decides who owns the card between renders. For a + phase-offloading pipeline (the default) activation is the ownership + claim and the transfers stay per-phase; only a pipeline loaded with + `phase_offload = FALSE` is copied to the card wholesale, and that + path is checked against free VRAM first. `resident_status()` reports + `components_on_gpu` alongside `state`, because the two legitimately + disagree: a render returns every component to pinned host memory as + its phase ends, so an active handle can hold nothing. + + Verified on an RTX 5060 Ti (16 GB) against local artifacts: flux2 + (11.22 GB pinned, 9.1 s render), flux1 (15.73 GB, 40.1 s), zimage + (13.45 GB, 19.5 s) and ltx (18.41 GB across 5 components). All three + image models reproduce bit-for-bit across a deactivate/activate + cycle. FLUX.1 and LTX both have pinned sets larger than the card, so + they exercise the refusal path rather than bulk onload. + # diffuseR 0.2.1.1 Addresses the CRAN review of the 0.2.0 submission. diff --git a/R/resident.R b/R/resident.R new file mode 100644 index 0000000..df7d14f --- /dev/null +++ b/R/resident.R @@ -0,0 +1,620 @@ +# In-process model residency: pinned host weights, disposable GPU copies. +# +# A resident pipeline keeps its canonical weights as page-locked (pinned) +# CPU tensors for the life of the handle. Activation creates the GPU +# representation with a DMA copy from pinned memory; deactivation +# destroys only the GPU representation and re-points the modules at the +# pinned host storage. Reactivation never touches the disk, so handing +# a small GPU between models is a sub-second operation instead of a full +# pipeline reload. +# +# Same contract as whisper's and chatterbox's R/resident.R (the three +# packages present one interface to a residency broker), adapted to +# diffuseR's shape: a pipeline is a classed list holding SEVERAL +# nn_modules (transformer, decoder, text encoder(s), and for LTX a video +# VAE, audio VAE and vocoder), so components are discovered by scanning +# the pipeline for nn_module fields and the staging set is keyed by +# field name. Non-tensor fields (config, tokenizer, scheduler) ride +# along untouched. +# +# This layer sits ABOVE the per-generation phase offloading already in +# txt2img_flux()/txt2img_flux2()/txt2img_zimage()/txt2vid_ltx23(). Those +# swap one component at a time WITHIN a render; residency is about who +# owns the GPU BETWEEN renders. The two compose: an active handle renders +# with its normal internal phase behaviour, and deactivation releases +# whatever is still resident so a sibling model can take the card. +# +# Mechanics rest on two torch behaviours (verified in whisper's port, +# re-verified by this package's tests): +# - nn_module$to() REBINDS parameter/buffer objects, so pinned host +# tensors held in res$staging survive activation, and any tensor handle +# taken before a transition is stale after it. All re-binding therefore +# resolves the modules' CURRENT tensors by name, every time. +# - Tensor$set_data() works across devices: a CUDA parameter can be +# re-pointed directly at a pinned CPU tensor. That is the evict +# mechanism; the orphaned CUDA storage is reclaimed by gc() + +# cuda_empty_cache(). +# +# States: inactive -> activating -> active -> deactivating -> inactive. +# Failed transitions roll back to pinned host state; a rollback that +# cannot be verified leaves the handle "broken" (fail-closed: only status +# and unload work). "unloaded" is terminal. + +# Families that ship a pinned/staged loader. Keyed by the `model` name +# used everywhere else in the package (see recommend()). +.resident_families <- c("flux1", "flux2", "zimage", "ltx") + +#' Every nn_module field of a pipeline, by name +#' +#' Discovery beats a hard-coded list: the families disagree on which +#' components exist (FLUX.1 has two text encoders, LTX adds a video VAE, +#' an audio VAE and a vocoder), and a field added later is picked up +#' without touching this file. +#' +#' @param pipeline A loaded diffuseR pipeline. +#' +#' @return A named list of the pipeline's \code{nn_module} fields, +#' possibly empty. +#' +#' @keywords internal +.resident_components <- function(pipeline) { + keep <- vapply(pipeline, function(x) inherits(x, "nn_module"), logical(1)) + pipeline[keep] +} + +#' Pin every component of a pipeline for fast transfer +#' +#' Re-uses any staging the loader already built (the phase-offload path +#' pins as part of loading), and pins the rest. Pinning a component that +#' is currently on the GPU also evicts it, since \code{.pin_component} +#' copies into page-locked host memory and re-points the live tensors at +#' it, so this doubles as the initial offload. +#' +#' @param pipeline A loaded diffuseR pipeline. +#' @param verbose Print progress. +#' +#' @return A named list of staging sets, one per component that could be +#' pinned. Components that fail to page-lock are absent, and fall back +#' to the pageable \code{$to()} path. +#' +#' @keywords internal +.resident_pin <- function(pipeline, verbose = TRUE) { + existing <- pipeline$staging %||% list() + comps <- .resident_components(pipeline) + if (verbose && length(comps)) { + message("Pinning ", length(comps), " components for residency...") + } + staging <- list() + for (nm in names(comps)) { + if (!is.null(existing[[nm]])) { + staging[[nm]] <- existing[[nm]] + next + } + extra <- if (identical(nm, "transformer") && + isTRUE(pipeline$fp8_resident)) { + .flux_fp8_collect(comps[[nm]]) + } else { + NULL + } + st <- .pin_component(comps[[nm]], extra = extra) + if (!is.null(st)) { + staging[[nm]] <- st + } + } + staging +} + +#' Total pinned host bytes across a staging set +#' +#' @param staging A named list of staging sets. +#' +#' @return Numeric. Bytes of page-locked host memory held. +#' +#' @keywords internal +.resident_pinned_bytes <- function(staging) { + total <- 0 + for (st in staging) { + for (pair in st) { + total <- total + prod(as.numeric(pair$pinned$shape)) * + .dtype_bytes(pair$pinned$dtype) + } + } + total +} + +# Bytes per element, keyed by the libtorch dtype name that +# as.character() on a torch_dtype returns ("Float", "Half", "Byte", +# "Long", ...), NOT the R constructor alias. Unknown dtypes fall back to +# 4, which only affects a reported number. +.dtype_widths <- c(double = 8, long = 8, complexfloat = 8, + float = 4, int = 4, + half = 2, bfloat16 = 2, short = 2, + byte = 1, char = 1, bool = 1, + float8_e4m3fn = 1, float8_e5m2 = 1) + +.dtype_bytes <- function(dtype) { + nm <- tolower(tryCatch(as.character(dtype), error = function(e) "")) + w <- .dtype_widths[[nm, exact = TRUE]] + if (is.null(w)) 4 else w +} + +#' TRUE when every staged tensor sits on the expected device type +#' +#' @param staging A named list of staging sets. +#' @param type "cpu" or "cuda". +#' +#' @return Logical. +#' +#' @keywords internal +.resident_all_on <- function(staging, type) { + for (st in staging) { + for (pair in st) { + dev <- tryCatch(pair$live$device$type, error = function(e) NA_character_) + if (!identical(dev, type)) { + return(FALSE) + } + } + } + TRUE +} + +#' How many components actually have their tensors on the GPU +#' +#' Ground truth, as opposed to the handle's declared state. The two can +#' disagree: a pipeline built with \code{phase_offload = TRUE} swaps each +#' component back to pinned host memory as its phase finishes, so after a +#' render the handle is still "active" while the card holds nothing. A +#' broker deciding who to evict needs the measurement, not the claim. +#' +#' @param staging A named list of staging sets. +#' +#' @return Integer. Number of components whose live tensors are on CUDA. +#' +#' @keywords internal +.resident_on_gpu_count <- function(staging) { + sum(vapply(staging, function(st) { + isTRUE(length(st) > 0) && + identical(tryCatch(st[[1]]$live$device$type, + error = function(e) NA_character_), "cuda") + }, logical(1))) +} + +#' Refuse operations that the current state cannot serve +#' +#' @param res A resident handle. +#' @param verb What the caller is attempting, for the message. +#' +#' @return Invisibly TRUE, or an error. +#' +#' @keywords internal +.resident_guard <- function(res, verb) { + if (identical(res$state, "unloaded")) { + stop("cannot ", verb, ": this handle is unloaded", call. = FALSE) + } + if (identical(res$state, "broken")) { + stop("cannot ", verb, ": this handle is broken (", res$last_error %||% + "no detail recorded", "). Only resident_status() and ", + "resident_unload() work from here.", call. = FALSE) + } + invisible(TRUE) +} + +#' Load a diffusion pipeline as a resident handle +#' +#' Loads a pipeline once and keeps its weights page-locked on the host +#' for the life of the handle. The GPU representation is created by +#' \code{\link{resident_activate}} and destroyed by +#' \code{\link{resident_deactivate}}, so a 16 GB card can hand itself +#' between models without either one re-reading its weights from disk. +#' +#' The handle is bound to one explicit GPU at load: a bare \code{"cuda"} +#' resolves to the current device now, and every later transition uses +#' that index, so the handle cannot drift to whichever GPU happens to be +#' current at transition time. +#' +#' One caveat on multi-GPU hosts: the family loader itself runs on the +#' \emph{current} device, and only the residency handle is bound to +#' \code{device}. Loading with \code{device = "cuda:1"} from a session +#' whose current device is 0 therefore stages through GPU 0 before the +#' first activation lands on GPU 1. Wrap the call in +#' \code{torch::with_device(device = "cuda:1", ...)} when that matters. +#' +#' The pipeline is left \emph{inactive} (weights pinned on the host, no +#' VRAM held). Call \code{\link{resident_activate}} before generating. +#' +#' @param model One of "flux1", "flux2", "zimage", "ltx". +#' @param device Target CUDA device, e.g. "cuda" or "cuda:1". +#' @param ... Passed to the family loader (\code{\link{flux_load_pipeline}}, +#' \code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}} +#' or \code{\link{ltx23_load_pipeline}}). \code{ltx} requires +#' \code{checkpoint_path}. +#' @param verbose Print progress messages. +#' +#' @return A \code{diffuseR_resident} handle (an environment). Inspect it +#' with \code{\link{resident_status}}; the fields of interest are the +#' state, the bound device, the component names, and the pinned host +#' byte count. +#' +#' @seealso \code{\link{resident_activate}}, \code{\link{resident_status}} +#' +#' @examples +#' \dontrun{ +#' res <- resident_load("flux2") +#' resident_activate(res) +#' img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +#' resident_deactivate(res) # VRAM freed, weights stay pinned in RAM +#' resident_activate(res) # fast: DMA copy, no disk +#' resident_unload(res) +#' } +#' +#' @export +resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), + device = "cuda", ..., verbose = TRUE) { + model <- match.arg(model) + if (!torch::cuda_is_available()) { + stop("resident_load() requires CUDA", call. = FALSE) + } + if (!grepl("^cuda", device)) { + stop("resident_load() requires a CUDA device, got '", device, "'", + call. = FALSE) + } + # Bind to one explicit GPU now, so later transitions cannot drift. + bound <- if (identical(device, "cuda")) { + paste0("cuda:", torch::cuda_current_device()) + } else { + device + } + + loader <- switch(model, + flux1 = flux_load_pipeline, + flux2 = flux2_load_pipeline, + zimage = zimage_load_pipeline, + ltx = ltx23_load_pipeline) + # Capture the phase-offload choice here rather than reading it back + # off the pipeline: the FLUX family stores it as a field, LTX takes + # it again at generate time and stores nothing, so the field is + # absent there and a NULL would be misread as "stays resident". + # Every family loader defaults it TRUE. + dots <- list(...) + phase_offload <- if (is.null(dots$phase_offload)) { + TRUE + } else { + isTRUE(dots$phase_offload) + } + pipeline <- loader(device = "cuda", verbose = verbose, ...) + + staging <- .resident_pin(pipeline, verbose = verbose) + # Pinning also evicted anything the loader had left resident, so the + # handle starts inactive with no VRAM held. + .resident_release_vram() + + res <- new.env(parent = emptyenv()) + res$model <- model + res$device <- bound + res$pipeline <- pipeline + res$phase_offload <- phase_offload + res$staging <- staging + res$components <- names(.resident_components(pipeline)) + res$pinned_bytes <- .resident_pinned_bytes(staging) + res$state <- "inactive" + res$last_error <- NULL + res$loaded_at <- Sys.time() + structure(res, class = "diffuseR_resident") +} + +#' Refuse a bulk activation that cannot fit +#' +#' Fails before the transfer rather than part-way through it. A partial +#' onload that OOMs is recoverable (activation rolls back), but it wastes +#' the transfer and reports a libtorch allocator error instead of the +#' actual problem, which is that this model does not fit this card. +#' +#' @param res A resident handle. +#' +#' @return Invisibly TRUE, or an error naming both figures. +#' +#' @keywords internal +.resident_check_fits <- function(res) { + free_gb <- tryCatch(.detect_vram(use_free = TRUE), + error = function(e) NA_real_) + need_gb <- res$pinned_bytes / 1024^3 + if (!is.na(free_gb) && free_gb > 0 && need_gb > free_gb) { + stop(sprintf(paste0("%s needs %.2f GB resident but only %.2f GB of ", + "VRAM is free. Load the pipeline with ", + "phase_offload = TRUE so components move on ", + "one phase at a time."), + res$model, need_gb, free_gb), call. = FALSE) + } + invisible(TRUE) +} + +# gc() then empty the caching allocator. Split out so every transition +# releases VRAM the same way. +.resident_release_vram <- function() { + gc() + tryCatch(torch::cuda_empty_cache(), error = function(e) NULL) + invisible(NULL) +} + +#' Bring a resident pipeline onto the GPU +#' +#' Copies every pinned component to the handle's bound device by DMA and +#' verifies the result tensor-by-tensor. A failure rolls back to the +#' pinned host state; a rollback that cannot itself be verified leaves +#' the handle broken. +#' +#' What activation does depends on how the pipeline was loaded: +#' +#' \itemize{ +#' \item \code{phase_offload = TRUE} (the default, and what the +#' \code{txt2img_*} functions expect): no bulk transfer. The render +#' moves each component on as its phase begins and back off as it +#' ends, from these same pinned copies, so pre-loading them would be +#' undone within one phase. Activation is the ownership claim. +#' \item \code{phase_offload = FALSE}: every component is copied to the +#' card up front and stays there across renders. This is the fast +#' path, and it is checked against free VRAM first. +#' } +#' +#' The distinction is not cosmetic. FLUX.1's pinned set is 15.73 GB, +#' which does not fit a 15.47 GiB card -- bulk-onloading it OOMs even +#' though the phased render fits comfortably. So \code{state} is a claim +#' about who owns the card, not a measurement of what is on it; read +#' \code{components_on_gpu} from \code{\link{resident_status}} for the +#' measurement. +#' +#' @param res A \code{diffuseR_resident} handle. +#' +#' @return Invisibly the handle, with state "active". +#' +#' @export +resident_activate <- function(res) { + stopifnot(inherits(res, "diffuseR_resident")) + .resident_guard(res, "activate") + if (identical(res$state, "active")) { + return(invisible(res)) + } + if (!identical(res$state, "inactive")) { + stop("cannot activate from state '", res$state, "'", call. = FALSE) + } + res$state <- "activating" + # A phase-offloading pipeline moves each component onto the card as + # its phase begins and straight back off as it ends, from these same + # pinned copies. Bulk-onloading here is therefore redundant -- the + # render undoes it within one phase -- and actively harmful: FLUX.1's + # pinned set is 15.73 GB, which does not fit a 15.47 GiB card even + # though the phased render does. For those pipelines, activation is + # the ownership claim; the transfers stay per-phase. + # The loader's own field wins when it kept one (the FLUX family + # downgrades phase_offload to FALSE on a CPU device); otherwise fall + # back to what resident_load() was asked for. + bulk <- !isTRUE(res$pipeline$phase_offload %||% res$phase_offload) + ok <- tryCatch({ + if (bulk) { + .resident_check_fits(res) + for (nm in names(res$staging)) { + .staged_onload(res$staging[[nm]], res$device) + } + } + TRUE + }, error = function(e) { + res$last_error <- conditionMessage(e) + FALSE + }) + if (ok) { + res$state <- "active" + return(invisible(res)) + } + # Roll back to pinned host state and verify it. + rolled <- tryCatch({ + for (nm in names(res$staging)) { + .staged_offload(res$staging[[nm]]) + } + .resident_release_vram() + .resident_all_on(res$staging, "cpu") + }, error = function(e) FALSE) + if (isTRUE(rolled)) { + res$state <- "inactive" + stop("resident_activate() failed (rolled back to pinned host ", + "state): ", res$last_error, call. = FALSE) + } + res$state <- "broken" + stop("resident_activate() failed and the rollback could not be ", + "verified: ", res$last_error, call. = FALSE) +} + +#' Release a resident pipeline's VRAM +#' +#' Re-points every component at its pinned host copy and drops the GPU +#' storage. Weights are immutable during inference, so the pinned copies +#' are still current and this moves no bytes: it is a pointer swap plus a +#' cache release. The handle stays loaded and can be reactivated without +#' touching the disk. +#' +#' @param res A \code{diffuseR_resident} handle. +#' @param release Empty the CUDA caching allocator afterwards. Leave TRUE +#' unless another handle on the same device is about to reuse the pool. +#' +#' @return Invisibly the handle, with state "inactive". +#' +#' @export +resident_deactivate <- function(res, release = TRUE) { + stopifnot(inherits(res, "diffuseR_resident")) + .resident_guard(res, "deactivate") + if (identical(res$state, "inactive")) { + return(invisible(res)) + } + if (!identical(res$state, "active")) { + stop("cannot deactivate from state '", res$state, "'", call. = FALSE) + } + res$state <- "deactivating" + verified <- tryCatch({ + for (nm in names(res$staging)) { + .staged_offload(res$staging[[nm]]) + } + if (isTRUE(release)) { + .resident_release_vram() + } + .resident_all_on(res$staging, "cpu") + }, error = function(e) { + res$last_error <- conditionMessage(e) + FALSE + }) + if (isTRUE(verified)) { + res$state <- "inactive" + return(invisible(res)) + } + res$state <- "broken" + stop("resident_deactivate() could not verify the pinned host state; ", + "the handle is broken and holds no usable GPU copy. ", + res$last_error %||% "", call. = FALSE) +} + +#' Generate from an active resident pipeline +#' +#' Dispatches to the family's generator with the resident pipeline +#' supplied, so no weights are re-read. The handle must be active. +#' +#' @param res A \code{diffuseR_resident} handle. +#' @param prompt Character. The text prompt. +#' @param ... Passed to \code{\link{txt2img_flux}}, +#' \code{\link{txt2img_flux2}}, \code{\link{txt2img_zimage}} or +#' \code{\link{txt2vid_ltx2}}. +#' +#' @return Whatever the family generator returns: an image array for the +#' image families, a video array for \code{ltx}. +#' +#' @export +resident_generate <- function(res, prompt, ...) { + stopifnot(inherits(res, "diffuseR_resident")) + .resident_guard(res, "generate") + if (!identical(res$state, "active")) { + stop("cannot generate from state '", res$state, + "'; call resident_activate() first", call. = FALSE) + } + gen <- switch(res$model, + flux1 = txt2img_flux, + flux2 = txt2img_flux2, + zimage = txt2img_zimage, + ltx = txt2vid_ltx2) + gen(prompt, pipeline = res$pipeline, ...) +} + +#' Status of a resident handle +#' +#' @param res A \code{diffuseR_resident} handle. +#' +#' @return A list with \code{model}, \code{state}, \code{device}, +#' \code{components} (character vector of pinned component names), +#' \code{pinned_bytes} (page-locked host bytes held), +#' \code{gpu_allocated} and \code{gpu_reserved} (bytes the CUDA +#' caching allocator reports live and held for this process, NA +#' without CUDA), \code{components_on_gpu} (how many components are +#' *actually* resident right now), \code{loaded_at}, and +#' \code{last_error} (NULL unless a transition failed). +#' +#' \code{state} is the handle's claim on the card; +#' \code{components_on_gpu} is the measurement. They disagree by +#' design after a render on a \code{phase_offload = TRUE} pipeline, +#' which returns each component to pinned host memory as its phase +#' finishes: the handle stays "active" (it still owns the card's +#' budget and can render again without touching disk) while +#' \code{components_on_gpu} is 0. Schedule on the measurement. +#' +#' @export +resident_status <- function(res) { + stopifnot(inherits(res, "diffuseR_resident")) + mem <- .cuda_bytes() + list(model = res$model, + state = res$state, + device = res$device, + components = names(res$staging), + components_on_gpu = .resident_on_gpu_count(res$staging), + pinned_bytes = res$pinned_bytes, + gpu_allocated = mem$allocated, + gpu_reserved = mem$reserved, + loaded_at = res$loaded_at, + last_error = res$last_error) +} + +# Live and reserved CUDA bytes for this process. torch has no +# cuda_memory_allocated(); the numbers live under cuda_memory_stats(), +# which itself errors without a CUDA build, hence the tryCatch. +.cuda_bytes <- function() { + s <- tryCatch(torch::cuda_memory_stats(), error = function(e) NULL) + if (is.null(s)) { + return(list(allocated = NA_real_, reserved = NA_real_)) + } + list(allocated = s$allocated_bytes$all$current %||% NA_real_, + reserved = s$reserved_bytes$all$current %||% NA_real_) +} + +#' Drop a resident handle entirely +#' +#' Releases the GPU copy if any, drops the pipeline and the pinned host +#' storage, and marks the handle unloaded. Terminal: nothing but +#' \code{\link{resident_status}} works afterwards. +#' +#' @param res A \code{diffuseR_resident} handle. +#' +#' @return Invisibly the handle, with state "unloaded". +#' +#' @export +resident_unload <- function(res) { + stopifnot(inherits(res, "diffuseR_resident")) + if (identical(res$state, "unloaded")) { + return(invisible(res)) + } + # Best effort: a broken handle still gets its memory back. + tryCatch({ + for (nm in names(res$staging)) { + .staged_offload(res$staging[[nm]]) + } + }, error = function(e) NULL) + res$pipeline <- NULL + res$staging <- list() + res$components <- character(0) + res$pinned_bytes <- 0 + res$state <- "unloaded" + .resident_release_vram() + invisible(res) +} + +#' Print a resident handle +#' +#' @param x A \code{diffuseR_resident} handle. +#' @param ... Ignored. +#' +#' @return Invisibly \code{x}. Called for the side effect of printing a +#' one-block summary to the console. +#' +#' @export +print.diffuseR_resident <- function(x, ...) { + s <- resident_status(x) + cat("\n") + cat(" model: ", s$model, "\n", sep = "") + cat(" state: ", s$state, "\n", sep = "") + cat(" device: ", s$device, "\n", sep = "") + cat(" components: ", + if (length(s$components)) paste(s$components, collapse = ", ") else "-", + "\n", sep = "") + cat(" on gpu: ", s$components_on_gpu, " of ", length(s$components), + "\n", sep = "") + cat(" pinned: ", .fmt_gb(s$pinned_bytes), "\n", sep = "") + if (!is.na(s$gpu_allocated)) { + cat(" gpu: ", .fmt_gb(s$gpu_allocated), " allocated, ", + .fmt_gb(s$gpu_reserved), " reserved\n", sep = "") + } + if (!is.null(s$last_error)) { + cat(" last error: ", s$last_error, "\n", sep = "") + } + invisible(x) +} + +# Byte count as GB, for the print method and messages. +.fmt_gb <- function(b) { + if (is.null(b) || is.na(b) || b <= 0) { + return("0 GB") + } + sprintf("%.2f GB", b / 1024^3) +} diff --git a/inst/tinytest/test_resident.R b/inst/tinytest/test_resident.R new file mode 100644 index 0000000..59890a6 --- /dev/null +++ b/inst/tinytest/test_resident.R @@ -0,0 +1,187 @@ +# Residency contract: component discovery, byte accounting, the state +# machine, and the pinned<->GPU round trip. +# +# The state-machine tests build a handle by hand around synthetic +# nn_modules, so they exercise the transitions without loading a real +# multi-GB pipeline. The round-trip test needs CUDA and is skipped +# without it. + +library(tinytest) +library(diffuseR) + +fake_pipeline <- function() { + structure(list(transformer = torch::nn_linear(8, 8), + decoder = torch::nn_linear(8, 4), + text_encoder = torch::nn_linear(4, 8), + config = list(family = "test"), + scheduler = "flowmatch", + phase_offload = TRUE), + class = "test_pipeline") +} + +# --- component discovery ---------------------------------------------------------- + +pipe <- fake_pipeline() +comps <- diffuseR:::.resident_components(pipe) +expect_equal(sort(names(comps)), c("decoder", "text_encoder", "transformer")) +# Non-module fields are not components. +expect_false("config" %in% names(comps)) +expect_false("scheduler" %in% names(comps)) +expect_false("phase_offload" %in% names(comps)) + +# A pipeline with no modules yields an empty set, not an error. +expect_equal(length(diffuseR:::.resident_components(list(a = 1, b = "x"))), 0L) + +# --- dtype byte table ------------------------------------------------------------- + +expect_equal(diffuseR:::.dtype_bytes(torch::torch_float32()), 4) +expect_equal(diffuseR:::.dtype_bytes(torch::torch_float16()), 2) +expect_equal(diffuseR:::.dtype_bytes(torch::torch_bfloat16()), 2) +expect_equal(diffuseR:::.dtype_bytes(torch::torch_uint8()), 1) +expect_equal(diffuseR:::.dtype_bytes(torch::torch_int64()), 8) + +# --- byte formatting -------------------------------------------------------------- + +expect_equal(diffuseR:::.fmt_gb(0), "0 GB") +expect_equal(diffuseR:::.fmt_gb(NA), "0 GB") +expect_equal(diffuseR:::.fmt_gb(NULL), "0 GB") +expect_equal(diffuseR:::.fmt_gb(1024^3), "1.00 GB") + +# --- state guard ------------------------------------------------------------------ + +mk <- function(state, staging = list(), phase_offload = FALSE, + pinned_bytes = 0, store_field = TRUE) { + e <- new.env(parent = emptyenv()) + e$model <- "flux2" + e$device <- "cuda:0" + e$state <- state + e$staging <- staging + e$components <- character(0) + e$pinned_bytes <- pinned_bytes + e$last_error <- NULL + e$loaded_at <- Sys.time() + # store_field = FALSE mimics LTX, which never records phase_offload + # on the pipeline object. + e$pipeline <- if (store_field) list(phase_offload = phase_offload) else list() + e$phase_offload <- phase_offload + structure(e, class = "diffuseR_resident") +} + +# A pipeline that does not store the field must still be read as +# phase-offloading (this is what sent LTX down the bulk path and made it +# demand 18.41 GB on a 15 GB card). +expect_true(isTRUE(mk("inactive", phase_offload = TRUE, + store_field = FALSE)$phase_offload)) + +expect_error(diffuseR:::.resident_guard(mk("unloaded"), "activate"), + pattern = "unloaded") +expect_error(diffuseR:::.resident_guard(mk("broken"), "activate"), + pattern = "broken") +expect_true(diffuseR:::.resident_guard(mk("inactive"), "activate")) +expect_true(diffuseR:::.resident_guard(mk("active"), "generate")) + +# --- illegal transitions ---------------------------------------------------------- + +# Deactivating something that was never activated is an error, not a +# silent no-op, unless it is already inactive. +expect_silent(resident_deactivate(mk("inactive"))) +expect_error(resident_deactivate(mk("activating")), pattern = "cannot deactivate") +expect_error(resident_activate(mk("deactivating")), pattern = "cannot activate") + +# Activate on an already-active handle is idempotent. +expect_silent(resident_activate(mk("active"))) + +# Generation is refused unless active. +expect_error(resident_generate(mk("inactive"), "a cat"), + pattern = "resident_activate") +expect_error(resident_generate(mk("broken"), "a cat"), pattern = "broken") + +# --- a set that cannot fit is refused before any transfer ------------------------- + +# Named figures beat a libtorch allocator error: FLUX.1's 15.73 GB pinned +# set does not fit a 15.47 GiB card, and the message should say so. +huge <- mk("inactive", pinned_bytes = 1e15) +expect_error(diffuseR:::.resident_check_fits(huge), pattern = "needs") +expect_error(diffuseR:::.resident_check_fits(huge), pattern = "phase_offload") +# A trivial set is never refused. +expect_true(diffuseR:::.resident_check_fits(mk("inactive", pinned_bytes = 0))) + +# --- status and print ------------------------------------------------------------- + +h <- mk("inactive") +s <- resident_status(h) +expect_true(all(c("model", "state", "device", "components", + "components_on_gpu", "pinned_bytes", "gpu_allocated", + "gpu_reserved", "loaded_at", "last_error") %in% names(s))) +# An empty staging set has nothing resident. +expect_equal(s$components_on_gpu, 0L) +# The allocator numbers must come from a torch function that actually +# exists: cuda_memory_allocated() does not (R CMD check catches it, but +# only as a WARNING buried in the dependencies step). +expect_true("cuda_memory_stats" %in% getNamespaceExports("torch")) +mem <- diffuseR:::.cuda_bytes() +expect_true(all(c("allocated", "reserved") %in% names(mem))) +expect_equal(s$state, "inactive") +expect_null(s$last_error) + +out <- capture.output(print(h)) +expect_true(any(grepl("diffuseR resident", out))) +expect_true(any(grepl("flux2", out))) + +# --- unload is terminal and idempotent -------------------------------------------- + +u <- mk("inactive") +resident_unload(u) +expect_equal(u$state, "unloaded") +expect_silent(resident_unload(u)) +expect_error(resident_activate(u), pattern = "unloaded") +# Status still works on an unloaded handle. +expect_equal(resident_status(u)$state, "unloaded") + +# --- CUDA round trip -------------------------------------------------------------- + +if (at_home() && torch::cuda_is_available()) { + pipe <- fake_pipeline() + staging <- diffuseR:::.resident_pin(pipe, verbose = FALSE) + expect_equal(sort(names(staging)), + c("decoder", "text_encoder", "transformer")) + + # Pinning leaves everything on the host. + expect_true(diffuseR:::.resident_all_on(staging, "cpu")) + expect_true(diffuseR:::.resident_pinned_bytes(staging) > 0) + + res <- mk("inactive", staging) + res$components <- names(staging) + + resident_activate(res) + expect_equal(res$state, "active") + expect_true(diffuseR:::.resident_all_on(staging, "cuda")) + expect_equal(resident_status(res)$components_on_gpu, 3L) + + resident_deactivate(res) + expect_equal(res$state, "inactive") + expect_true(diffuseR:::.resident_all_on(staging, "cpu")) + expect_equal(resident_status(res)$components_on_gpu, 0L) + + # components_on_gpu is ground truth, not the declared state: evict + # behind the handle's back and status must notice while state still + # says "active". This is what a phase-offloading render does. + resident_activate(res) + for (nm in names(staging)) diffuseR:::.staged_offload(staging[[nm]]) + expect_equal(res$state, "active") + expect_equal(resident_status(res)$components_on_gpu, 0L) + resident_deactivate(res) + + # A second round trip reuses the same pinned buffers. + resident_activate(res) + expect_true(diffuseR:::.resident_all_on(staging, "cuda")) + resident_deactivate(res) + expect_true(diffuseR:::.resident_all_on(staging, "cpu")) + + # The module still computes correctly after the round trip. + y <- torch::with_no_grad(pipe$transformer(torch::torch_randn(c(2, 8)))) + expect_equal(as.integer(y$shape), c(2L, 8L)) + + resident_unload(res) + expect_equal(res$state, "unloaded") +} diff --git a/man/dot-resident_all_on.Rd b/man/dot-resident_all_on.Rd new file mode 100644 index 0000000..6b7a8eb --- /dev/null +++ b/man/dot-resident_all_on.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_all_on} +\alias{.resident_all_on} +\title{TRUE when every staged tensor sits on the expected device type} +\usage{ +.resident_all_on(staging, type) +} +\arguments{ +\item{staging}{A named list of staging sets.} + +\item{type}{"cpu" or "cuda".} +} +\value{ +Logical. +} +\description{ +TRUE when every staged tensor sits on the expected device type +} +\keyword{internal} diff --git a/man/dot-resident_check_fits.Rd b/man/dot-resident_check_fits.Rd new file mode 100644 index 0000000..46cdec9 --- /dev/null +++ b/man/dot-resident_check_fits.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_check_fits} +\alias{.resident_check_fits} +\title{Refuse a bulk activation that cannot fit} +\usage{ +.resident_check_fits(res) +} +\arguments{ +\item{res}{A resident handle.} +} +\value{ +Invisibly TRUE, or an error naming both figures. +} +\description{ +Fails before the transfer rather than part-way through it. A partial +onload that OOMs is recoverable (activation rolls back), but it wastes +the transfer and reports a libtorch allocator error instead of the +actual problem, which is that this model does not fit this card. +} +\keyword{internal} diff --git a/man/dot-resident_components.Rd b/man/dot-resident_components.Rd new file mode 100644 index 0000000..c8be784 --- /dev/null +++ b/man/dot-resident_components.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_components} +\alias{.resident_components} +\title{Every nn_module field of a pipeline, by name} +\usage{ +.resident_components(pipeline) +} +\arguments{ +\item{pipeline}{A loaded diffuseR pipeline.} +} +\value{ +A named list of the pipeline's \code{nn_module} fields, + possibly empty. +} +\description{ +Discovery beats a hard-coded list: the families disagree on which +components exist (FLUX.1 has two text encoders, LTX adds a video VAE, +an audio VAE and a vocoder), and a field added later is picked up +without touching this file. +} +\keyword{internal} diff --git a/man/dot-resident_guard.Rd b/man/dot-resident_guard.Rd new file mode 100644 index 0000000..d9d5562 --- /dev/null +++ b/man/dot-resident_guard.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_guard} +\alias{.resident_guard} +\title{Refuse operations that the current state cannot serve} +\usage{ +.resident_guard(res, verb) +} +\arguments{ +\item{res}{A resident handle.} + +\item{verb}{What the caller is attempting, for the message.} +} +\value{ +Invisibly TRUE, or an error. +} +\description{ +Refuse operations that the current state cannot serve +} +\keyword{internal} diff --git a/man/dot-resident_on_gpu_count.Rd b/man/dot-resident_on_gpu_count.Rd new file mode 100644 index 0000000..108d4e4 --- /dev/null +++ b/man/dot-resident_on_gpu_count.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_on_gpu_count} +\alias{.resident_on_gpu_count} +\title{How many components actually have their tensors on the GPU} +\usage{ +.resident_on_gpu_count(staging) +} +\arguments{ +\item{staging}{A named list of staging sets.} +} +\value{ +Integer. Number of components whose live tensors are on CUDA. +} +\description{ +Ground truth, as opposed to the handle's declared state. The two can +disagree: a pipeline built with \code{phase_offload = TRUE} swaps each +component back to pinned host memory as its phase finishes, so after a +render the handle is still "active" while the card holds nothing. A +broker deciding who to evict needs the measurement, not the claim. +} +\keyword{internal} diff --git a/man/dot-resident_pin.Rd b/man/dot-resident_pin.Rd new file mode 100644 index 0000000..dd1583b --- /dev/null +++ b/man/dot-resident_pin.Rd @@ -0,0 +1,25 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_pin} +\alias{.resident_pin} +\title{Pin every component of a pipeline for fast transfer} +\usage{ +.resident_pin(pipeline, verbose = TRUE) +} +\arguments{ +\item{pipeline}{A loaded diffuseR pipeline.} + +\item{verbose}{Print progress.} +} +\value{ +A named list of staging sets, one per component that could be + pinned. Components that fail to page-lock are absent, and fall back + to the pageable \code{$to()} path. +} +\description{ +Re-uses any staging the loader already built (the phase-offload path +pins as part of loading), and pins the rest. Pinning a component that +is currently on the GPU also evicts it, since \code{.pin_component} +copies into page-locked host memory and re-points the live tensors at +it, so this doubles as the initial offload. +} +\keyword{internal} diff --git a/man/dot-resident_pinned_bytes.Rd b/man/dot-resident_pinned_bytes.Rd new file mode 100644 index 0000000..d8cd722 --- /dev/null +++ b/man/dot-resident_pinned_bytes.Rd @@ -0,0 +1,17 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_pinned_bytes} +\alias{.resident_pinned_bytes} +\title{Total pinned host bytes across a staging set} +\usage{ +.resident_pinned_bytes(staging) +} +\arguments{ +\item{staging}{A named list of staging sets.} +} +\value{ +Numeric. Bytes of page-locked host memory held. +} +\description{ +Total pinned host bytes across a staging set +} +\keyword{internal} diff --git a/man/print.diffuseR_resident.Rd b/man/print.diffuseR_resident.Rd new file mode 100644 index 0000000..5bf285d --- /dev/null +++ b/man/print.diffuseR_resident.Rd @@ -0,0 +1,19 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{print.diffuseR_resident} +\alias{print.diffuseR_resident} +\title{Print a resident handle} +\usage{ +\method{print}{diffuseR_resident}(x, ...) +} +\arguments{ +\item{x}{A \code{diffuseR_resident} handle.} + +\item{...}{Ignored.} +} +\value{ +Invisibly \code{x}. Called for the side effect of printing a + one-block summary to the console. +} +\description{ +Print a resident handle +} diff --git a/man/resident_activate.Rd b/man/resident_activate.Rd new file mode 100644 index 0000000..1c14087 --- /dev/null +++ b/man/resident_activate.Rd @@ -0,0 +1,41 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_activate} +\alias{resident_activate} +\title{Bring a resident pipeline onto the GPU} +\usage{ +resident_activate(res) +} +\arguments{ +\item{res}{A \code{diffuseR_resident} handle.} +} +\value{ +Invisibly the handle, with state "active". +} +\description{ +Copies every pinned component to the handle's bound device by DMA and +verifies the result tensor-by-tensor. A failure rolls back to the +pinned host state; a rollback that cannot itself be verified leaves +the handle broken. +} +\details{ +What activation does depends on how the pipeline was loaded: + +\itemize{ +\item \code{phase_offload = TRUE} (the default, and what the +\code{txt2img_*} functions expect): no bulk transfer. The render +moves each component on as its phase begins and back off as it +ends, from these same pinned copies, so pre-loading them would be +undone within one phase. Activation is the ownership claim. +\item \code{phase_offload = FALSE}: every component is copied to the +card up front and stays there across renders. This is the fast +path, and it is checked against free VRAM first. +} + +The distinction is not cosmetic. FLUX.1's pinned set is 15.73 GB, +which does not fit a 15.47 GiB card -- bulk-onloading it OOMs even +though the phased render fits comfortably. So \code{state} is a claim +about who owns the card, not a measurement of what is on it; read +\code{components_on_gpu} from \code{\link{resident_status}} for the +measurement. + +} diff --git a/man/resident_deactivate.Rd b/man/resident_deactivate.Rd new file mode 100644 index 0000000..07d3d36 --- /dev/null +++ b/man/resident_deactivate.Rd @@ -0,0 +1,23 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_deactivate} +\alias{resident_deactivate} +\title{Release a resident pipeline's VRAM} +\usage{ +resident_deactivate(res, release = TRUE) +} +\arguments{ +\item{res}{A \code{diffuseR_resident} handle.} + +\item{release}{Empty the CUDA caching allocator afterwards. Leave TRUE +unless another handle on the same device is about to reuse the pool.} +} +\value{ +Invisibly the handle, with state "inactive". +} +\description{ +Re-points every component at its pinned host copy and drops the GPU +storage. Weights are immutable during inference, so the pinned copies +are still current and this moves no bytes: it is a pointer swap plus a +cache release. The handle stays loaded and can be reactivated without +touching the disk. +} diff --git a/man/resident_generate.Rd b/man/resident_generate.Rd new file mode 100644 index 0000000..3eedd1d --- /dev/null +++ b/man/resident_generate.Rd @@ -0,0 +1,24 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_generate} +\alias{resident_generate} +\title{Generate from an active resident pipeline} +\usage{ +resident_generate(res, prompt, ...) +} +\arguments{ +\item{res}{A \code{diffuseR_resident} handle.} + +\item{prompt}{Character. The text prompt.} + +\item{...}{Passed to \code{\link{txt2img_flux}}, +\code{\link{txt2img_flux2}}, \code{\link{txt2img_zimage}} or +\code{\link{txt2vid_ltx2}}.} +} +\value{ +Whatever the family generator returns: an image array for the + image families, a video array for \code{ltx}. +} +\description{ +Dispatches to the family's generator with the resident pipeline +supplied, so no weights are re-read. The handle must be active. +} diff --git a/man/resident_load.Rd b/man/resident_load.Rd new file mode 100644 index 0000000..da4bb2d --- /dev/null +++ b/man/resident_load.Rd @@ -0,0 +1,64 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_load} +\alias{resident_load} +\title{Load a diffusion pipeline as a resident handle} +\usage{ +resident_load(model = c("flux2", "flux1", "zimage", "ltx"), device = "cuda", + ..., verbose = TRUE) +} +\arguments{ +\item{model}{One of "flux1", "flux2", "zimage", "ltx".} + +\item{device}{Target CUDA device, e.g. "cuda" or "cuda:1".} + +\item{...}{Passed to the family loader (\code{\link{flux_load_pipeline}}, +\code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}} +or \code{\link{ltx23_load_pipeline}}). \code{ltx} requires +\code{checkpoint_path}.} + +\item{verbose}{Print progress messages.} +} +\value{ +A \code{diffuseR_resident} handle (an environment). Inspect it + with \code{\link{resident_status}}; the fields of interest are the + state, the bound device, the component names, and the pinned host + byte count. +} +\description{ +Loads a pipeline once and keeps its weights page-locked on the host +for the life of the handle. The GPU representation is created by +\code{\link{resident_activate}} and destroyed by +\code{\link{resident_deactivate}}, so a 16 GB card can hand itself +between models without either one re-reading its weights from disk. +} +\details{ +The handle is bound to one explicit GPU at load: a bare \code{"cuda"} +resolves to the current device now, and every later transition uses +that index, so the handle cannot drift to whichever GPU happens to be +current at transition time. + +One caveat on multi-GPU hosts: the family loader itself runs on the +\emph{current} device, and only the residency handle is bound to +\code{device}. Loading with \code{device = "cuda:1"} from a session +whose current device is 0 therefore stages through GPU 0 before the +first activation lands on GPU 1. Wrap the call in +\code{torch::with_device(device = "cuda:1", ...)} when that matters. + +The pipeline is left \emph{inactive} (weights pinned on the host, no +VRAM held). Call \code{\link{resident_activate}} before generating. + +} +\examples{ +\dontrun{ +res <- resident_load("flux2") +resident_activate(res) +img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +resident_deactivate(res) # VRAM freed, weights stay pinned in RAM +resident_activate(res) # fast: DMA copy, no disk +resident_unload(res) +} + +} +\seealso{ +\code{\link{resident_activate}}, \code{\link{resident_status}} +} diff --git a/man/resident_status.Rd b/man/resident_status.Rd new file mode 100644 index 0000000..49a6d10 --- /dev/null +++ b/man/resident_status.Rd @@ -0,0 +1,31 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_status} +\alias{resident_status} +\title{Status of a resident handle} +\usage{ +resident_status(res) +} +\arguments{ +\item{res}{A \code{diffuseR_resident} handle.} +} +\value{ +A list with \code{model}, \code{state}, \code{device}, + \code{components} (character vector of pinned component names), + \code{pinned_bytes} (page-locked host bytes held), + \code{gpu_allocated} and \code{gpu_reserved} (bytes the CUDA + caching allocator reports live and held for this process, NA + without CUDA), \code{components_on_gpu} (how many components are + *actually* resident right now), \code{loaded_at}, and + \code{last_error} (NULL unless a transition failed). + + \code{state} is the handle's claim on the card; + \code{components_on_gpu} is the measurement. They disagree by + design after a render on a \code{phase_offload = TRUE} pipeline, + which returns each component to pinned host memory as its phase + finishes: the handle stays "active" (it still owns the card's + budget and can render again without touching disk) while + \code{components_on_gpu} is 0. Schedule on the measurement. +} +\description{ +Status of a resident handle +} diff --git a/man/resident_unload.Rd b/man/resident_unload.Rd new file mode 100644 index 0000000..25eddce --- /dev/null +++ b/man/resident_unload.Rd @@ -0,0 +1,18 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{resident_unload} +\alias{resident_unload} +\title{Drop a resident handle entirely} +\usage{ +resident_unload(res) +} +\arguments{ +\item{res}{A \code{diffuseR_resident} handle.} +} +\value{ +Invisibly the handle, with state "unloaded". +} +\description{ +Releases the GPU copy if any, drops the pipeline and the pinned host +storage, and marks the handle unloaded. Terminal: nothing but +\code{\link{resident_status}} works afterwards. +}