diff --git a/DESCRIPTION b/DESCRIPTION index ae7adbb..09b61ac 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.3 +Version: 0.2.2.4 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 abb78e0..7e28492 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -193,6 +193,7 @@ export(save_video) export(save_video_ltx23) export(scheduler_add_noise) export(sd_pipeline_from_safetensors) +export(sdxl_load_pipeline) export(sdxl_memory_profile) export(sdxl_pipeline_from_safetensors) export(serve) diff --git a/NEWS.md b/NEWS.md index 6eea834..573c309 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,29 @@ +# diffuseR 0.2.2.4 + +* `resident_load()` accepts `"sdxl"`, making it the fifth resident family. + `sdxl_load_pipeline()` is the new adapter: it defaults to the + `download_sdxl()` cache, fixes the UNet at float16 before the weights are + pinned, and marks the pipeline so that only the UNet is placed on the + card. Onloading all four components fits the 8.0 GB of weights and then + OOMs in the VAE decode, which runs 1024x1024 in float32 while the UNet is + still resident; the text encode and decode therefore run on the host from + the same pinned copies. `resident_generate()` supplies the matching + `devices` so `txt2img_sdxl()` does not re-decide the placement with + `auto_devices()`, and an explicit `devices` from the caller still wins. + +* `resident_deactivate()` now releases the NF4 dequantisation buffers for + `ltx`. They live in a package-level environment rather than in the module, + so offloading the weights did not free them and `gc()` could not reclaim + them, leaving scratch on the card after every deactivation. + +* A bulk activation pre-warms the CUDA caching allocator. Growing the pool + one allocation per tensor dominated the first activation: 24.16 s against + 0.32 s once warm. + +* `txt2img_sdxl()` no longer requires the legacy TorchScript `.pt` files + when it is handed a pipeline it did not build, and `setup_dtype()` accepts + an ordinal-qualified device such as `"cuda:0"`. + # diffuseR 0.2.2.3 * Prebuilt NF4 artifacts for flux2 (2.1 GB) and zimage (3.5 GB) are now diff --git a/R/models2devices.R b/R/models2devices.R index e0f6bb8..65cdb84 100644 --- a/R/models2devices.R +++ b/R/models2devices.R @@ -1,3 +1,34 @@ +#' Device configuration for an already-built pipeline +#' +#' The half of \code{\link{models2devices}} that does not touch the disk. +#' +#' \code{models2devices()} ends by calling \code{download_model()}, which +#' resolves the TorchScript \code{.pt} files for the model and stops with +#' "Missing model files" when they are absent. That is correct when it is +#' about to load them, and wrong when the caller already holds a pipeline: +#' a native safetensors pipeline never reads a \code{.pt}, so verifying them +#' makes a working generation depend on files it does not use. Passing +#' \code{download_models = FALSE} does not avoid it -- the check runs either +#' way and only the downloading is suppressed. +#' +#' So callers with a pipeline in hand take this path and get the same four +#' fields without the file check. +#' +#' @param model_name A character string naming the model, e.g. "sdxl". +#' @param devices A device string or named list of component devices. +#' @param unet_dtype_str A character string naming the UNet dtype, or NULL. +#' +#' @return The same shape \code{\link{models2devices}} returns: +#' \code{devices}, \code{unet_dtype}, \code{device_cpu}, \code{device_cuda}. +#' +#' @keywords internal +.devices_for_pipeline <- function(model_name, devices, unet_dtype_str = NULL) { + dv <- standardize_devices(devices, get_required_components(model_name)) + list(devices = dv, unet_dtype = setup_dtype(dv, unet_dtype_str), + device_cpu = torch::torch_device("cpu"), + device_cuda = torch::torch_device("cuda")) +} + #' models2devices #' @description This function sets up the model directory, device configuration, and data types for diffusion models. #' It checks the validity of the model name and devices, detects model type, and downloads the model if necessary. @@ -129,6 +160,12 @@ setup_dtype <- function(devices, unet_dtype_str) { } else { stop("No main computation component found") } + # An ordinal-qualified device picks the same dtype as the bare one. + # resident_load() binds to an explicit "cuda:N" so later transitions + # cannot drift, and passes that through as the component device; without + # this, "cuda:0" matches neither branch below and falls to the "Invalid + # device" stop. + main_device <- sub(":.*$", "", main_device) if (main_device == "cpu") { return(torch::torch_float32()) diff --git a/R/resident.R b/R/resident.R index ed532e2..64f7c99 100644 --- a/R/resident.R +++ b/R/resident.R @@ -42,7 +42,7 @@ # 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") +.resident_families <- c("flux1", "flux2", "zimage", "ltx", "sdxl") #' Every nn_module field of a pipeline, by name #' @@ -226,12 +226,14 @@ #' 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 model One of "flux1", "flux2", "zimage", "ltx", "sdxl". #' @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}. +#' \code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}}, +#' \code{\link{ltx23_load_pipeline}} or +#' \code{\link{sdxl_load_pipeline}}). \code{ltx} requires +#' \code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the +#' \code{\link{download_sdxl}} cache). #' @param verbose Print progress messages. #' #' @return A \code{diffuseR_resident} handle (an environment). Inspect it @@ -252,7 +254,7 @@ #' } #' #' @export -resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), +resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", "sdxl"), device = "cuda", ..., verbose = TRUE) { model <- match.arg(model) if (!torch::cuda_is_available()) { @@ -273,7 +275,8 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), flux1 = flux_load_pipeline, flux2 = flux2_load_pipeline, zimage = zimage_load_pipeline, - ltx = ltx23_load_pipeline) + ltx = ltx23_load_pipeline, + sdxl = sdxl_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 @@ -298,6 +301,9 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), res$pipeline <- pipeline res$phase_offload <- phase_offload res$staging <- staging + # A family may want only part of its set on the card (see + # .resident_gpu_set). NULL means all of it. + res$gpu_components <- pipeline$gpu_components res$components <- names(.resident_components(pipeline)) res$pinned_bytes <- .resident_pinned_bytes(staging) res$state <- "inactive" @@ -306,6 +312,78 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), structure(res, class = "diffuseR_resident") } +#' Grow the caching allocator's pool in one allocation before a bulk onload +#' +#' A cold bulk onload grows the pool one \code{cudaMalloc} per tensor, and +#' the syscalls dominate: SDXL's 8.0 GB pinned set measured 24.16 s on the +#' first activation against 0.32 s on the second, a 74x ratio, on an +#' otherwise idle RTX 5060 Ti. One large allocation freed straight back into +#' the pool lets the transfers carve from cached blocks instead. Same +#' technique the NF4 LTX loader already uses for its first transformer +#' onload (83.5 s -> 4.6 s there). +#' +#' It matters beyond the wall clock: a residency broker with a startup +#' deadline reads 24 s inside a first activate as a wedged worker. +#' +#' Best-effort. A card that cannot seat the block in one piece falls back to +#' the per-tensor path, which is the current behaviour and merely slow, so +#' the failure is swallowed rather than raised. +#' +#' @param bytes Numeric. Host bytes about to be transferred; the pool is +#' warmed to this plus a small margin for allocator slack. +#' @param device Target CUDA device. +#' +#' @return Invisibly NULL. +#' +#' @keywords internal +.resident_prewarm <- function(bytes, device) { + # isTRUE() rather than a bare is.finite(): a NULL pinned_bytes gives + # logical(0), and `||` on a zero-length value is an error in R >= 4.3, + # so the guard meant to skip the pre-warm would instead fail the + # activation it exists to speed up. + if (!isTRUE(is.finite(bytes)) || bytes <= 0) { + return(invisible(NULL)) + } + tryCatch({ + warm <- torch::torch_empty(as.numeric(bytes) * 1.05, + dtype = torch::torch_uint8(), + device = device) + rm(warm) + gc(verbose = FALSE) + }, error = function(e) invisible(NULL)) + invisible(NULL) +} + +#' Which components a bulk activation puts on the card +#' +#' All of them, unless the family says otherwise. +#' +#' SDXL says otherwise. Its four components are only 8.0 GB pinned, so +#' bulk-onloading the set looks affordable on a 16 GB card -- and then the +#' VAE decode OOMs, because SDXL decodes 1024x1024 in float32 and that peak +#' arrives while the UNet is still resident. Measured: 8.0 GB of weights +#' plus the decode phase reached 14.38 GiB of 15.47 GiB and died asking for +#' another 512 MiB. A 12 GB card never had a chance. +#' +#' So SDXL puts only the UNet on the card and computes the text encode and +#' the decode on the host from the same pinned copies. That is the placement +#' \code{\link{auto_devices}} already recommends for this model at this tier; +#' residency's contribution is that the 5 GB UNet stops being re-read from +#' disk between models. +#' +#' @param res A resident handle. +#' +#' @return Character vector of \code{res$staging} names. +#' +#' @keywords internal +.resident_gpu_set <- function(res) { + want <- res$gpu_components + if (is.null(want)) { + return(names(res$staging)) + } + intersect(want, names(res$staging)) +} + #' Refuse a bulk activation that cannot fit #' #' Fails before the transfer rather than part-way through it. A partial @@ -319,16 +397,24 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx"), #' which means "cannot tell" and never refuses, so a test that wants #' the refusal has to state the budget rather than depend on the #' machine having a card. +#' @param need_bytes Bytes actually headed for the card. NULL means the +#' whole pinned set, which is right for a family that onloads everything +#' and wrong for one that onloads a subset -- SDXL pins 8.0 GB and sends +#' 5.1 GB of it, so charging it the full figure would refuse activations +#' that fit. #' #' @return Invisibly TRUE, or an error naming both figures. #' #' @keywords internal -.resident_check_fits <- function(res, free_gb = NULL) { +.resident_check_fits <- function(res, free_gb = NULL, need_bytes = NULL) { if (is.null(free_gb)) { free_gb <- tryCatch(.detect_vram(use_free = TRUE), error = function(e) NA_real_) } - need_gb <- res$pinned_bytes / 1024 ^ 3 + if (is.null(need_bytes)) { + need_bytes <- res$pinned_bytes + } + need_gb <- need_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 ", @@ -402,8 +488,11 @@ resident_activate <- function(res) { bulk <- !isTRUE(res$pipeline$phase_offload %||% res$phase_offload) ok <- tryCatch({ if (bulk) { - .resident_check_fits(res) - for (nm in names(res$staging)) { + onto <- .resident_gpu_set(res) + need <- .resident_pinned_bytes(res$staging[onto]) + .resident_check_fits(res, need_bytes = need) + .resident_prewarm(need, res$device) + for (nm in onto) { .staged_onload(res$staging[[nm]], res$device) } } @@ -463,6 +552,22 @@ resident_deactivate <- function(res, release = TRUE) { for (nm in names(res$staging)) { .staged_offload(res$staging[[nm]]) } + # NF4 linears dequantize into a package-level scratch environment, + # not into the module, so offloading the weights does not free it. + # txt2vid_ltx2() deliberately SKIPS its own release when the + # transformer is resident (the next chunk reuses the buffers), which + # is right within a render and leaves the scratch on the card once + # the render is over. Nothing else reclaims it: gc() and + # cuda_empty_cache() cannot touch a block the environment still + # references. Deactivation is the moment the handle stops owning the + # card, so it is the moment to drop them. Unconditional on `release` + # -- a broker that passes release = FALSE to keep the pool warm for + # the next tenant is exactly the caller that must not be handed a + # budget short by this scratch. The image families already release + # it at the end of every generate. + if (identical(res$model, "ltx")) { + ltx23_release_dequant_buffers() + } if (isTRUE(release)) { .resident_release_vram() } @@ -489,11 +594,18 @@ resident_deactivate <- function(res, release = TRUE) { #' @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}. +#' \code{\link{txt2img_flux2}}, \code{\link{txt2img_zimage}}, +#' \code{\link{txt2vid_ltx2}} or \code{\link{txt2img_sdxl}}. For +#' \code{sdxl} the handle supplies \code{devices} matching its own +#' placement unless the caller names it. +#' +#' @return Whatever the family generator returns, and the families do not +#' agree: an image array for \code{flux1}, \code{flux2} and +#' \code{zimage}, a video array for \code{ltx}, and for \code{sdxl} a +#' list of \code{image} and \code{metadata}, because +#' \code{\link{txt2img_sdxl}} has always returned that pair and changing +#' it would break every existing caller. A broker that wants one shape +#' should normalise in its own wrapper. #' #' @export resident_generate <- function(res, prompt, ...) { @@ -507,8 +619,51 @@ resident_generate <- function(res, prompt, ...) { flux1 = txt2img_flux, flux2 = txt2img_flux2, zimage = txt2img_zimage, - ltx = txt2vid_ltx2) - gen(prompt, pipeline = res$pipeline, ...) + ltx = txt2vid_ltx2, + sdxl = txt2img_sdxl) + do.call(gen, c(list(prompt, pipeline = res$pipeline), + .resident_gen_args(res, list(...)))) +} + +#' Family fixups for a resident generate call +#' +#' Split out so the SDXL device injection can be asserted without running a +#' multi-gigabyte generation. +#' +#' \code{\link{txt2img_sdxl}} does not read the pipeline's placement. With +#' its default \code{devices = "auto"} it calls \code{\link{auto_devices}} +#' afresh and moves the prompt embeds to whatever THAT returns. On a 12 GB +#' card auto can answer "unet on cuda, encoders on cpu", which contradicts a +#' bulk-activated handle whose encoders are on the card, and the text +#' encoder call then dies on a device mismatch. The handle knows where its +#' components actually are, so it says so instead of letting the generator +#' re-decide. +#' +#' Only SDXL needs this: the other families phase-offload from their own +#' pinned copies and place each component themselves as its phase begins. +#' +#' @param res A resident handle. +#' @param args The caller's \code{...}, as a list. An explicit +#' \code{devices} wins -- this fills a gap, it does not override. +#' +#' @return \code{args}, possibly with \code{devices} added. +#' +#' @keywords internal +.resident_gen_args <- function(res, args) { + if (identical(res$model, "sdxl") && is.null(args$devices)) { + on_gpu <- .resident_gpu_set(res) + place <- function(nm) { + if (nm %in% on_gpu) { + res$device + } else { + "cpu" + } + } + args$devices <- list(unet = place("unet"), decoder = place("decoder"), + text_encoder = place("text_encoder"), + text_encoder2 = place("text_encoder2")) + } + args } #' Status of a resident handle diff --git a/R/sdxl_pipeline_safetensors.R b/R/sdxl_pipeline_safetensors.R index a721175..7b294ec 100644 --- a/R/sdxl_pipeline_safetensors.R +++ b/R/sdxl_pipeline_safetensors.R @@ -126,6 +126,100 @@ sdxl_pipeline_from_safetensors <- function(diffusers_dir, devices = NULL, native_decode = TRUE) } +#' Load the SDXL pipeline in the family-loader convention +#' +#' The adapter \code{\link{resident_load}} needs. +#' \code{\link{sdxl_pipeline_from_safetensors}} predates the residency layer +#' and has its own signature: a required \code{diffusers_dir} and a plural +#' \code{devices} list, where every other family loader takes an optional +#' model directory and a singular \code{device}. This translates. +#' +#' Two choices are not cosmetic: +#' +#' \code{unet_dtype} is fixed HERE rather than at generation time. +#' \code{sdxl_pipeline_from_safetensors} defaults it from the component +#' device, so loading to CPU for pinning would page-lock a float32 UNet +#' (~10 GB) and then render in float32. A resident handle must decide the +#' dtype from where it will COMPUTE, not from where the weights are parked +#' while pinned. +#' +#' \code{phase_offload} is FALSE, unlike every other family. SDXL has no +#' per-phase offloading path -- \code{\link{txt2img_sdxl}} places components +#' once and leaves them -- so activation has to be a real transfer rather +#' than an ownership claim. \code{\link{resident_activate}} reads this field +#' off the pipeline and does the right thing. +#' +#' The pipeline also carries \code{gpu_components = "unet"}: all four +#' components are pinned, but only the UNet is put on the card, and the text +#' encode and VAE decode run on the host. The 8.0 GB pinned set makes +#' onloading everything look affordable on a 16 GB card, and it is not -- +#' SDXL decodes 1024x1024 in float32 and that peak lands while the UNet is +#' still resident, which reached 14.38 GiB of 15.47 GiB and OOMed. On the +#' 12 GB cards this wrapper exists for, only the UNet was ever going to fit. +#' +#' @param model_dir Diffusers directory (with \code{unet/}, \code{vae/}, +#' \code{text_encoder/}, \code{text_encoder_2/}). NULL, the default, +#' resolves the \code{\link{download_sdxl}} cache, fetching it if absent. +#' @param device Where the pipeline will COMPUTE once activated. Components +#' are built on the CPU regardless, because residency pins them there and +#' \code{\link{resident_activate}} moves them; this only picks the dtype. +#' @param unet_dtype A torch dtype for the UNet. NULL picks float16 for a +#' CUDA device and float32 for CPU. +#' @param phase_offload Kept for signature parity with the other family +#' loaders. SDXL has no phased path, so anything but FALSE is ignored. +#' @param verbose Logical. +#' +#' @return The list from \code{\link{sdxl_pipeline_from_safetensors}}, plus +#' \code{phase_offload}. +#' +#' @seealso \code{\link{resident_load}} +#' +#' @examples +#' \dontrun{ +#' res <- resident_load("sdxl") +#' resident_activate(res) +#' img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +#' resident_deactivate(res) +#' } +#' +#' @export +sdxl_load_pipeline <- function(model_dir = NULL, device = "cuda", + unet_dtype = NULL, phase_offload = FALSE, + verbose = TRUE) { + if (is.null(model_dir)) { + model_dir <- download_sdxl(verbose = verbose) + } + if (is.null(unet_dtype)) { + unet_dtype <- if (grepl("^cuda", device)) { + torch::torch_float16() + } else { + torch::torch_float32() + } + } + # All-CPU: .resident_pin() page-locks from here and activation moves the + # copies onto the card. Building straight onto the GPU would allocate a + # device copy that pinning immediately evicts. + pipeline <- sdxl_pipeline_from_safetensors( + model_dir, + devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu", + text_encoder2 = "cpu"), + unet_dtype = unet_dtype, verbose = verbose) + pipeline$phase_offload <- isTRUE(phase_offload) + # Only the UNet goes to the card. All four are pinned and the text + # encode and VAE decode run on the host from those same pinned copies. + # + # Bulk-onloading all four is what the 8.0 GB pinned figure invites, and + # it OOMs: SDXL decodes 1024x1024 in float32, and that peak arrives + # while the UNet is still resident. Measured on a 15.47 GiB card -- + # weights plus decode reached 14.38 GiB and died asking for another + # 512 MiB. A 12 GB card, which is the reason this wrapper exists, has + # no chance at all. txt2img_sdxl() has no way to evict the UNet before + # decode (its "phase cleanup" is gc + empty_cache, which cannot move a + # live module), so the decision has to be made here. + pipeline$gpu_components <- "unet" + pipeline +} + # Hosted on the cornball-ai/sdxl-R dataset under diffusers/. fp16, sub-2 GB # per file (the 5 GB UNet is re-sharded via reshard_safetensors so it is # CRAN-safetensors readable). The native UNet constructor uses the SDXL diff --git a/R/txt2img_sdxl.R b/R/txt2img_sdxl.R index 06e27b7..ddc3c43 100644 --- a/R/txt2img_sdxl.R +++ b/R/txt2img_sdxl.R @@ -99,8 +99,16 @@ txt2img_sdxl <- function(prompt, negative_prompt = NULL, img_dim = 1024, devices <- auto_devices(model_name) } - m2d <- models2devices(model_name = model_name, devices = devices, - unet_dtype_str = unet_dtype_str) + # A supplied pipeline needs the device/dtype resolution but not the + # TorchScript file check models2devices() ends with: a native pipeline + # reads no .pt, and verifying them makes an otherwise working call fail + # on files it never opens. + m2d <- if (is.null(pipeline)) { + models2devices(model_name = model_name, devices = devices, + unet_dtype_str = unet_dtype_str) + } else { + .devices_for_pipeline(model_name, devices, unet_dtype_str) + } devices <- m2d$devices unet_dtype <- m2d$unet_dtype device_cpu <- m2d$device_cpu diff --git a/inst/tinytest/test_resident_sdxl.R b/inst/tinytest/test_resident_sdxl.R new file mode 100644 index 0000000..15945f8 --- /dev/null +++ b/inst/tinytest/test_resident_sdxl.R @@ -0,0 +1,211 @@ +# SDXL residency, and the two fixes it needed on the way in. +# +# None of this requires a GPU or the 7 GB of weights: the staging layer is +# already covered by test_resident.R and test_staging.R, and what is new +# here is dispatch, the device/dtype resolution that a supplied pipeline +# takes, and the NF4 scratch release on deactivate. + +library(tinytest) +library(diffuseR) + +# --- dispatch --------------------------------------------------------------------- + +expect_true("sdxl" %in% diffuseR:::.resident_families) + +# The match.arg default has to list it too, or resident_load("sdxl") +# is refused before any of the above matters. Read off the formals rather +# than restated, so the two cannot drift apart. +expect_true("sdxl" %in% eval(formals(resident_load)$model)) +expect_equal(sort(eval(formals(resident_load)$model)), + sort(diffuseR:::.resident_families)) + +expect_true(is.function(sdxl_load_pipeline)) + +# --- the loader shim --------------------------------------------------------------- + +fm <- formals(sdxl_load_pipeline) +expect_true(all(c("model_dir", "device", "unet_dtype", "phase_offload", + "verbose") %in% names(fm))) + +# model_dir defaults to NULL so the download_sdxl() cache is resolved. +expect_null(fm$model_dir) + +# SDXL has no phased path, so a resident handle must transfer on activate +# rather than treat activation as a claim. This default is what +# resident_activate() reads to decide. +expect_false(fm$phase_offload) + +# --- the fit check charges only what is actually going to the card ------------------ + +# SDXL pins 8.0 GB and sends 5.1 GB of it. Charging the full pinned figure +# would refuse activations that fit. +h <- new.env(parent = emptyenv()) +h$model <- "sdxl" +h$pinned_bytes <- 8 * 1024^3 +expect_error(diffuseR:::.resident_check_fits(structure(h, class = "diffuseR_resident"), + free_gb = 6), + pattern = "needs 8") +expect_true(diffuseR:::.resident_check_fits(structure(h, class = "diffuseR_resident"), + free_gb = 6, + need_bytes = 5.1 * 1024^3)) + +# --- dtype resolution --------------------------------------------------------------- + +# resident_load() binds an explicit "cuda:N" so later transitions cannot +# drift, and passes it through as the component device. Before the fix that +# matched neither branch of setup_dtype() and hit "Invalid device". +expect_equal(diffuseR:::setup_dtype(list(unet = "cuda:0"), NULL), + torch::torch_float16()) +expect_equal(diffuseR:::setup_dtype(list(unet = "cuda:3"), NULL), + torch::torch_float16()) +# Unqualified devices keep their existing meaning. +expect_equal(diffuseR:::setup_dtype(list(unet = "cuda"), NULL), + torch::torch_float16()) +expect_equal(diffuseR:::setup_dtype(list(unet = "cpu"), NULL), + torch::torch_float32()) +# An explicit dtype still wins over the device's default. +expect_equal(diffuseR:::setup_dtype(list(unet = "cuda:0"), "float32"), + torch::torch_float32()) +# A genuinely unknown device is still refused. +expect_error(diffuseR:::setup_dtype(list(unet = "tpu"), NULL), + pattern = "Invalid device") + +# --- device config without the file check -------------------------------------------- + +# models2devices() ends by verifying TorchScript .pt files. A native +# safetensors pipeline never reads one, so a caller holding a pipeline takes +# the path that skips it. Same four fields, no disk. +d <- diffuseR:::.devices_for_pipeline("sdxl", "cuda:0", NULL) +expect_equal(sort(names(d)), + c("device_cpu", "device_cuda", "devices", "unet_dtype")) +expect_equal(d$unet_dtype, torch::torch_float16()) +expect_true(all(c("unet", "decoder", "text_encoder", "text_encoder2") %in% + names(d$devices))) +expect_equal(d$devices$unet, "cuda:0") + +# A named list is carried through, not flattened to one device. +d2 <- diffuseR:::.devices_for_pipeline( + "sdxl", list(unet = "cuda", decoder = "cpu", + text_encoder = "cpu", text_encoder2 = "cpu"), NULL) +expect_equal(d2$devices$unet, "cuda") +expect_equal(d2$devices$decoder, "cpu") +# dtype follows the UNet, not the other components. +expect_equal(d2$unet_dtype, torch::torch_float16()) + +# --- generate-time device injection --------------------------------------------------- + +sdxl_names <- c("unet", "decoder", "text_encoder", "text_encoder2") + +mk_h <- function(model, device = "cuda:0", gpu = "unet", + comps = sdxl_names) { + e <- new.env(parent = emptyenv()) + e$model <- model + e$device <- device + e$gpu_components <- gpu + # .resident_gpu_set() intersects against the staging set, so the fake + # handle needs one; the values are never read. + e$staging <- stats::setNames(vector("list", length(comps)), comps) + structure(e, class = "diffuseR_resident") +} + +# --- which components go to the card --------------------------------------------- + +# NULL means everything, which is what the other families want. +expect_equal(sort(diffuseR:::.resident_gpu_set(mk_h("flux2", gpu = NULL))), + sort(sdxl_names)) +expect_equal(diffuseR:::.resident_gpu_set(mk_h("sdxl")), "unet") +# A named component that was never pinned is dropped rather than onloaded. +expect_equal(diffuseR:::.resident_gpu_set(mk_h("sdxl", gpu = c("unet", "nope"))), + "unet") + +inj <- diffuseR:::.resident_gen_args(mk_h("sdxl"), list()) +expect_true(!is.null(inj$devices)) +expect_equal(sort(names(inj$devices)), + c("decoder", "text_encoder", "text_encoder2", "unet")) + +# Only the UNet is on the card: bulk-onloading all four fits the weights and +# then OOMs in the fp32 VAE decode. The encoders and decoder compute on the +# host from their pinned copies. +expect_equal(inj$devices$unet, "cuda:0") +expect_equal(inj$devices$decoder, "cpu") +expect_equal(inj$devices$text_encoder, "cpu") +expect_equal(inj$devices$text_encoder2, "cpu") + +# The bound ordinal is carried, so a handle on the second card does not +# quietly render on the first. +expect_equal(diffuseR:::.resident_gen_args(mk_h("sdxl", "cuda:1"), + list())$devices$unet, "cuda:1") + +# The placement follows gpu_components rather than being hard-coded, so a +# roomier card can be given the decoder too without touching this logic. +wide <- diffuseR:::.resident_gen_args( + mk_h("sdxl", gpu = c("unet", "decoder")), list()) +expect_equal(wide$devices$decoder, "cuda:0") +expect_equal(wide$devices$text_encoder, "cpu") + +# An explicit devices= from the caller is a decision, not a gap to fill. +keep <- list(devices = list(unet = "cpu")) +expect_equal(diffuseR:::.resident_gen_args(mk_h("sdxl"), keep)$devices, + list(unet = "cpu")) + +# Other arguments ride through untouched. +o <- diffuseR:::.resident_gen_args(mk_h("sdxl"), list(seed = 7L)) +expect_equal(o$seed, 7L) + +# The phase-offloading families place components themselves; injecting a +# device list for them would fight their own per-phase movement. +for (m in c("flux1", "flux2", "zimage", "ltx")) { + expect_null(diffuseR:::.resident_gen_args(mk_h(m), list())$devices) +} + +# --- NF4 dequant scratch is released on deactivate ------------------------------------- + +# NF4 linears dequantize into a package-level environment rather than into +# the module, so offloading the weights does not free it, and +# txt2vid_ltx2() deliberately skips its own release while the transformer is +# resident. Nothing else reclaims it: the environment still holds a +# reference, so gc() and cuda_empty_cache() cannot. +mk_d <- function(model) { + e <- new.env(parent = emptyenv()) + e$model <- model + e$device <- "cuda:0" + e$state <- "active" + e$staging <- list() + e$pipeline <- list() + e$last_error <- NULL + structure(e, class = "diffuseR_resident") +} + +buf <- diffuseR:::.ltx23_dequant_buffers +assign("probe", 1L, envir = buf) +expect_true("probe" %in% ls(buf)) +resident_deactivate(mk_d("ltx"), release = FALSE) +expect_equal(length(ls(buf)), 0L) + +# Unconditional on `release`: a broker that passes release = FALSE to keep +# the pool warm for the next tenant is precisely the caller that must not +# be handed a budget short by this scratch. +assign("probe2", 1L, envir = buf) +resident_deactivate(mk_d("ltx"), release = TRUE) +expect_equal(length(ls(buf)), 0L) + +# Only LTX skips the release inside its own generate, so only LTX needs it +# here. The image families already clear it at the end of every render, and +# clearing it for them would be reaching into another family's business. +assign("probe3", 1L, envir = buf) +resident_deactivate(mk_d("flux2"), release = FALSE) +expect_true("probe3" %in% ls(buf)) +rm("probe3", envir = buf) + +# --- allocator pre-warm ---------------------------------------------------------------- + +# Best-effort by contract: it is an optimisation, and a card that cannot +# seat the block in one piece must fall back to the per-tensor path rather +# than fail the activation. +expect_silent(diffuseR:::.resident_prewarm(0, "cuda")) +expect_silent(diffuseR:::.resident_prewarm(-1, "cuda")) +expect_silent(diffuseR:::.resident_prewarm(NA_real_, "cuda")) +expect_silent(diffuseR:::.resident_prewarm(NULL, "cuda")) +# An impossible size on a real card must be swallowed, not raised. +expect_silent(diffuseR:::.resident_prewarm(1e18, "cuda")) +expect_null(diffuseR:::.resident_prewarm(1024, "cuda")) diff --git a/man/dot-devices_for_pipeline.Rd b/man/dot-devices_for_pipeline.Rd new file mode 100644 index 0000000..1bab063 --- /dev/null +++ b/man/dot-devices_for_pipeline.Rd @@ -0,0 +1,36 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.devices_for_pipeline} +\alias{.devices_for_pipeline} +\title{Device configuration for an already-built pipeline} +\usage{ +.devices_for_pipeline(model_name, devices, unet_dtype_str = NULL) +} +\arguments{ +\item{model_name}{A character string naming the model, e.g. "sdxl".} + +\item{devices}{A device string or named list of component devices.} + +\item{unet_dtype_str}{A character string naming the UNet dtype, or NULL.} +} +\value{ +The same shape \code{\link{models2devices}} returns: + \code{devices}, \code{unet_dtype}, \code{device_cpu}, \code{device_cuda}. +} +\description{ +The half of \code{\link{models2devices}} that does not touch the disk. +} +\details{ +\code{models2devices()} ends by calling \code{download_model()}, which +resolves the TorchScript \code{.pt} files for the model and stops with +"Missing model files" when they are absent. That is correct when it is +about to load them, and wrong when the caller already holds a pipeline: +a native safetensors pipeline never reads a \code{.pt}, so verifying them +makes a working generation depend on files it does not use. Passing +\code{download_models = FALSE} does not avoid it -- the check runs either +way and only the downloading is suppressed. + +So callers with a pipeline in hand take this path and get the same four +fields without the file check. + +} +\keyword{internal} diff --git a/man/dot-resident_check_fits.Rd b/man/dot-resident_check_fits.Rd index d698812..ad60c9c 100644 --- a/man/dot-resident_check_fits.Rd +++ b/man/dot-resident_check_fits.Rd @@ -3,7 +3,7 @@ \alias{.resident_check_fits} \title{Refuse a bulk activation that cannot fit} \usage{ -.resident_check_fits(res, free_gb = NULL) +.resident_check_fits(res, free_gb = NULL, need_bytes = NULL) } \arguments{ \item{res}{A resident handle.} @@ -13,6 +13,12 @@ make the decision deterministic: with no GPU the measurement is 0, which means "cannot tell" and never refuses, so a test that wants the refusal has to state the budget rather than depend on the machine having a card.} + +\item{need_bytes}{Bytes actually headed for the card. NULL means the +whole pinned set, which is right for a family that onloads everything +and wrong for one that onloads a subset -- SDXL pins 8.0 GB and sends +5.1 GB of it, so charging it the full figure would refuse activations +that fit.} } \value{ Invisibly TRUE, or an error naming both figures. diff --git a/man/dot-resident_gen_args.Rd b/man/dot-resident_gen_args.Rd new file mode 100644 index 0000000..3a1c165 --- /dev/null +++ b/man/dot-resident_gen_args.Rd @@ -0,0 +1,35 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_gen_args} +\alias{.resident_gen_args} +\title{Family fixups for a resident generate call} +\usage{ +.resident_gen_args(res, args) +} +\arguments{ +\item{res}{A resident handle.} + +\item{args}{The caller's \code{...}, as a list. An explicit +\code{devices} wins -- this fills a gap, it does not override.} +} +\value{ +\code{args}, possibly with \code{devices} added. +} +\description{ +Split out so the SDXL device injection can be asserted without running a +multi-gigabyte generation. +} +\details{ +\code{\link{txt2img_sdxl}} does not read the pipeline's placement. With +its default \code{devices = "auto"} it calls \code{\link{auto_devices}} +afresh and moves the prompt embeds to whatever THAT returns. On a 12 GB +card auto can answer "unet on cuda, encoders on cpu", which contradicts a +bulk-activated handle whose encoders are on the card, and the text +encoder call then dies on a device mismatch. The handle knows where its +components actually are, so it says so instead of letting the generator +re-decide. + +Only SDXL needs this: the other families phase-offload from their own +pinned copies and place each component themselves as its phase begins. + +} +\keyword{internal} diff --git a/man/dot-resident_gpu_set.Rd b/man/dot-resident_gpu_set.Rd new file mode 100644 index 0000000..746da87 --- /dev/null +++ b/man/dot-resident_gpu_set.Rd @@ -0,0 +1,32 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_gpu_set} +\alias{.resident_gpu_set} +\title{Which components a bulk activation puts on the card} +\usage{ +.resident_gpu_set(res) +} +\arguments{ +\item{res}{A resident handle.} +} +\value{ +Character vector of \code{res$staging} names. +} +\description{ +All of them, unless the family says otherwise. +} +\details{ +SDXL says otherwise. Its four components are only 8.0 GB pinned, so +bulk-onloading the set looks affordable on a 16 GB card -- and then the +VAE decode OOMs, because SDXL decodes 1024x1024 in float32 and that peak +arrives while the UNet is still resident. Measured: 8.0 GB of weights +plus the decode phase reached 14.38 GiB of 15.47 GiB and died asking for +another 512 MiB. A 12 GB card never had a chance. + +So SDXL puts only the UNet on the card and computes the text encode and +the decode on the host from the same pinned copies. That is the placement +\code{\link{auto_devices}} already recommends for this model at this tier; +residency's contribution is that the 5 GB UNet stops being re-read from +disk between models. + +} +\keyword{internal} diff --git a/man/dot-resident_prewarm.Rd b/man/dot-resident_prewarm.Rd new file mode 100644 index 0000000..df25031 --- /dev/null +++ b/man/dot-resident_prewarm.Rd @@ -0,0 +1,35 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.resident_prewarm} +\alias{.resident_prewarm} +\title{Grow the caching allocator's pool in one allocation before a bulk onload} +\usage{ +.resident_prewarm(bytes, device) +} +\arguments{ +\item{bytes}{Numeric. Host bytes about to be transferred; the pool is +warmed to this plus a small margin for allocator slack.} + +\item{device}{Target CUDA device.} +} +\value{ +Invisibly NULL. +} +\description{ +A cold bulk onload grows the pool one \code{cudaMalloc} per tensor, and +the syscalls dominate: SDXL's 8.0 GB pinned set measured 24.16 s on the +first activation against 0.32 s on the second, a 74x ratio, on an +otherwise idle RTX 5060 Ti. One large allocation freed straight back into +the pool lets the transfers carve from cached blocks instead. Same +technique the NF4 LTX loader already uses for its first transformer +onload (83.5 s -> 4.6 s there). +} +\details{ +It matters beyond the wall clock: a residency broker with a startup +deadline reads 24 s inside a first activate as a wedged worker. + +Best-effort. A card that cannot seat the block in one piece falls back to +the per-tensor path, which is the current behaviour and merely slow, so +the failure is swallowed rather than raised. + +} +\keyword{internal} diff --git a/man/resident_generate.Rd b/man/resident_generate.Rd index 3eedd1d..cc7c557 100644 --- a/man/resident_generate.Rd +++ b/man/resident_generate.Rd @@ -11,12 +11,19 @@ resident_generate(res, prompt, ...) \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}}.} +\code{\link{txt2img_flux2}}, \code{\link{txt2img_zimage}}, +\code{\link{txt2vid_ltx2}} or \code{\link{txt2img_sdxl}}. For +\code{sdxl} the handle supplies \code{devices} matching its own +placement unless the caller names it.} } \value{ -Whatever the family generator returns: an image array for the - image families, a video array for \code{ltx}. +Whatever the family generator returns, and the families do not + agree: an image array for \code{flux1}, \code{flux2} and + \code{zimage}, a video array for \code{ltx}, and for \code{sdxl} a + list of \code{image} and \code{metadata}, because + \code{\link{txt2img_sdxl}} has always returned that pair and changing + it would break every existing caller. A broker that wants one shape + should normalise in its own wrapper. } \description{ Dispatches to the family's generator with the resident pipeline diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 0aa1902..daca7fe 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -4,21 +4,23 @@ \title{Load a diffusion pipeline as a resident handle} \usage{ resident_load( - model = c("flux2", "flux1", "zimage", "ltx"), + model = c("flux2", "flux1", "zimage", "ltx", "sdxl"), device = "cuda", ..., verbose = TRUE ) } \arguments{ -\item{model}{One of "flux1", "flux2", "zimage", "ltx".} +\item{model}{One of "flux1", "flux2", "zimage", "ltx", "sdxl".} \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}.} +\code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}}, +\code{\link{ltx23_load_pipeline}} or +\code{\link{sdxl_load_pipeline}}). \code{ltx} requires +\code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the +\code{\link{download_sdxl}} cache).} \item{verbose}{Print progress messages.} } diff --git a/man/sdxl_load_pipeline.Rd b/man/sdxl_load_pipeline.Rd new file mode 100644 index 0000000..a1fd192 --- /dev/null +++ b/man/sdxl_load_pipeline.Rd @@ -0,0 +1,78 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{sdxl_load_pipeline} +\alias{sdxl_load_pipeline} +\title{Load the SDXL pipeline in the family-loader convention} +\usage{ +sdxl_load_pipeline( + model_dir = NULL, + device = "cuda", + unet_dtype = NULL, + phase_offload = FALSE, + verbose = TRUE +) +} +\arguments{ +\item{model_dir}{Diffusers directory (with \code{unet/}, \code{vae/}, +\code{text_encoder/}, \code{text_encoder_2/}). NULL, the default, +resolves the \code{\link{download_sdxl}} cache, fetching it if absent.} + +\item{device}{Where the pipeline will COMPUTE once activated. Components +are built on the CPU regardless, because residency pins them there and +\code{\link{resident_activate}} moves them; this only picks the dtype.} + +\item{unet_dtype}{A torch dtype for the UNet. NULL picks float16 for a +CUDA device and float32 for CPU.} + +\item{phase_offload}{Kept for signature parity with the other family +loaders. SDXL has no phased path, so anything but FALSE is ignored.} + +\item{verbose}{Logical.} +} +\value{ +The list from \code{\link{sdxl_pipeline_from_safetensors}}, plus + \code{phase_offload}. +} +\description{ +The adapter \code{\link{resident_load}} needs. +\code{\link{sdxl_pipeline_from_safetensors}} predates the residency layer +and has its own signature: a required \code{diffusers_dir} and a plural +\code{devices} list, where every other family loader takes an optional +model directory and a singular \code{device}. This translates. +} +\details{ +Two choices are not cosmetic: + +\code{unet_dtype} is fixed HERE rather than at generation time. +\code{sdxl_pipeline_from_safetensors} defaults it from the component +device, so loading to CPU for pinning would page-lock a float32 UNet +(~10 GB) and then render in float32. A resident handle must decide the +dtype from where it will COMPUTE, not from where the weights are parked +while pinned. + +\code{phase_offload} is FALSE, unlike every other family. SDXL has no +per-phase offloading path -- \code{\link{txt2img_sdxl}} places components +once and leaves them -- so activation has to be a real transfer rather +than an ownership claim. \code{\link{resident_activate}} reads this field +off the pipeline and does the right thing. + +The pipeline also carries \code{gpu_components = "unet"}: all four +components are pinned, but only the UNet is put on the card, and the text +encode and VAE decode run on the host. The 8.0 GB pinned set makes +onloading everything look affordable on a 16 GB card, and it is not -- +SDXL decodes 1024x1024 in float32 and that peak lands while the UNet is +still resident, which reached 14.38 GiB of 15.47 GiB and OOMed. On the +12 GB cards this wrapper exists for, only the UNet was ever going to fit. + +} +\examples{ +\dontrun{ +res <- resident_load("sdxl") +resident_activate(res) +img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +resident_deactivate(res) +} + +} +\seealso{ +\code{\link{resident_load}} +}