diff --git a/DESCRIPTION b/DESCRIPTION index 09b61ac..b1f6bb5 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.4 +Version: 0.2.2.5 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 7e28492..3bb9f58 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(sd21_load_pipeline) export(sdxl_load_pipeline) export(sdxl_memory_profile) export(sdxl_pipeline_from_safetensors) diff --git a/NEWS.md b/NEWS.md index 573c309..ecc67ba 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,22 @@ +# diffuseR 0.2.2.5 + +* `resident_load()` accepts `"sd21"`, the sixth resident family. + `sd21_load_pipeline()` defaults to the `download_sd21()` cache and pins + the UNet at **float32**: SD 2.1's attention overflows in float16 and the + pipeline returns all-NaN rather than raising, so a float16 resident would + generate blank images with nothing to catch it. Only the UNet is placed + on the card, matching what `auto_devices("sd21")` already recommends. + +* `txt2img_sd21()` no longer requires the legacy TorchScript `.pt` files + when it is handed a pipeline it did not build, matching the same fix made + for `txt2img_sdxl()` in 0.2.2.4. + +* Known limitation, unchanged by this release but now measured: SD 2.1's + native float32 path does not fit its own 768x768 default on a 15.47 GiB + card. The denoise wants 11.35 GB with a further 3.1 GB of allocator + slack, with or without residency (a plain `txt2img_sd21()` OOMs at + 11.494 GB). 512x512 uses 6.807 GB and is comfortable. + # diffuseR 0.2.2.4 * `resident_load()` accepts `"sdxl"`, making it the fifth resident family. diff --git a/R/resident.R b/R/resident.R index 64f7c99..6011cf4 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", "sdxl") +.resident_families <- c("flux1", "flux2", "zimage", "ltx", "sdxl", "sd21") #' Every nn_module field of a pipeline, by name #' @@ -254,7 +254,8 @@ #' } #' #' @export -resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", "sdxl"), +resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", + "sdxl", "sd21"), device = "cuda", ..., verbose = TRUE) { model <- match.arg(model) if (!torch::cuda_is_available()) { @@ -276,7 +277,8 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", "sdxl"), flux2 = flux2_load_pipeline, zimage = zimage_load_pipeline, ltx = ltx23_load_pipeline, - sdxl = sdxl_load_pipeline) + sdxl = sdxl_load_pipeline, + sd21 = sd21_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 @@ -620,7 +622,8 @@ resident_generate <- function(res, prompt, ...) { flux2 = txt2img_flux2, zimage = txt2img_zimage, ltx = txt2vid_ltx2, - sdxl = txt2img_sdxl) + sdxl = txt2img_sdxl, + sd21 = txt2img_sd21) do.call(gen, c(list(prompt, pipeline = res$pipeline), .resident_gen_args(res, list(...)))) } @@ -650,18 +653,22 @@ resident_generate <- function(res, prompt, ...) { #' #' @keywords internal .resident_gen_args <- function(res, args) { - if (identical(res$model, "sdxl") && is.null(args$devices)) { + if (res$model %in% c("sdxl", "sd21") && is.null(args$devices)) { on_gpu <- .resident_gpu_set(res) - place <- function(nm) { - if (nm %in% on_gpu) { + # The component set is the generator's, not the pipeline's: SD 2.1 + # declares an `encoder` (the VAE encoder img2img needs) that the + # text-to-image pipeline never builds, and standardize_devices() + # would otherwise refuse the list as missing a required component. + # It follows the decoder, which is where it would live anyway. + want <- get_required_components(res$model) + args$devices <- stats::setNames(lapply(want, function(nm) { + src <- if (identical(nm, "encoder")) "decoder" else nm + if (src %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")) + }), want) } args } diff --git a/R/sd_pipeline_safetensors.R b/R/sd_pipeline_safetensors.R index 205259a..4e767e1 100644 --- a/R/sd_pipeline_safetensors.R +++ b/R/sd_pipeline_safetensors.R @@ -170,3 +170,81 @@ sd_pipeline_from_safetensors <- function(diffusers_dir, model_name = "sd21", list(unet = unet, decoder = decoder, text_encoder = text_encoder) } + +#' Load the SD 2.1 pipeline in the family-loader convention +#' +#' The adapter \code{\link{resident_load}} needs, the counterpart to +#' \code{\link{sdxl_load_pipeline}}. +#' \code{\link{sd_pipeline_from_safetensors}} takes a required +#' \code{diffusers_dir} and a plural \code{devices} list; every family +#' loader takes an optional model directory and a singular \code{device}. +#' +#' \strong{The UNet defaults to float32, not float16.} This is the opposite +#' of \code{\link{sdxl_load_pipeline}} and it is not a preference: SD 2.1's +#' attention overflows in float16 and the pipeline returns all-NaN, so a +#' float16 resident would load, activate, generate and hand back a blank +#' image without raising anything. \code{\link{txt2img_sd21}} already +#' defaults its native path to float32 for this reason; the resident loader +#' has to make the same choice, because it fixes the dtype before the +#' weights are pinned and nothing downstream can undo it. +#' +#' Only the UNet goes to the card, as for SDXL, but for a different reason. +#' SDXL is undone by its float32 VAE decode; SD 2.1 never reaches the +#' decode. The float32 requirement above makes the denoise itself expensive +#' -- measured 11.49 GB allocated plus 3.0 GB of allocator slack at the +#' 768x768 default, which does not fit a 15.47 GiB card with the weights +#' resident as well. That is not a residency artifact: a plain +#' \code{\link{txt2img_sd21}} with the same devices and dtype OOMs at +#' 11.494 GB, within 0.04 GB of the resident run. +#' +#' \code{\link{auto_devices}} already answers "unet on cuda, the rest on +#' cpu" for this model at every strategy, so residency is agreeing with the +#' package rather than inventing a placement. +#' +#' @param model_dir Diffusers directory (with \code{unet/}, \code{vae/}, +#' \code{text_encoder/}). NULL, the default, resolves the +#' \code{\link{download_sd21}} 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. +#' @param unet_dtype A torch dtype for the UNet. NULL means float32. Pass +#' \code{torch::torch_float16()} only if you have reason to believe the +#' overflow above no longer applies. +#' @param phase_offload Kept for signature parity with the other family +#' loaders. SD 2.1 has no phased path, so anything but FALSE is ignored. +#' @param verbose Logical. +#' +#' @return The list from \code{\link{sd_pipeline_from_safetensors}}, plus +#' \code{phase_offload}. +#' +#' @seealso \code{\link{resident_load}}, \code{\link{sdxl_load_pipeline}} +#' +#' @examples +#' \dontrun{ +#' res <- resident_load("sd21") +#' resident_activate(res) +#' img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +#' resident_deactivate(res) +#' } +#' +#' @export +sd21_load_pipeline <- function(model_dir = NULL, device = "cuda", + unet_dtype = NULL, phase_offload = FALSE, + verbose = TRUE) { + if (is.null(model_dir)) { + model_dir <- download_sd21(verbose = verbose) + } + if (is.null(unet_dtype)) { + # float32 on CUDA as well as CPU. See the note above: float16 is a + # silent all-NaN, not an error. + unet_dtype <- torch::torch_float32() + } + pipeline <- sd_pipeline_from_safetensors(model_dir, model_name = "sd21", + devices = list(unet = "cpu", decoder = "cpu", text_encoder = "cpu"), + unet_dtype = unet_dtype, verbose = verbose) + pipeline$phase_offload <- isTRUE(phase_offload) + # See the note above: the float32 denoise alone wants ~11.5 GB, so the + # weights cannot all sit on the card beside it. + pipeline$gpu_components <- "unet" + pipeline +} diff --git a/R/txt2img_sd21.R b/R/txt2img_sd21.R index bca2733..565c688 100644 --- a/R/txt2img_sd21.R +++ b/R/txt2img_sd21.R @@ -56,9 +56,13 @@ txt2img_sd21 <- function(prompt, negative_prompt = NULL, img_dim = 768, device_cpu <- torch::torch_device("cpu") device_cuda <- torch::torch_device("cuda") - if (!is.null(diffusers_dir)) { + if (!is.null(diffusers_dir) || !is.null(pipeline)) { # Native safetensors path: resolve devices/dtype with the pure # helpers, skipping the .pt model verification models2devices runs. + # A supplied pipeline takes this branch too -- it was built without + # a .pt and never reads one, so verifying them would fail a working + # call on files it does not open. That is how a resident handle + # arrives here. # SD 2.1 attention overflows in float16 (all-NaN output), so # default this path to float32 unless float16 is asked for. devices <- standardize_devices(devices, diff --git a/inst/tinytest/test_resident_sd21.R b/inst/tinytest/test_resident_sd21.R new file mode 100644 index 0000000..04b8ff7 --- /dev/null +++ b/inst/tinytest/test_resident_sd21.R @@ -0,0 +1,105 @@ +# SD 2.1 residency. The sibling of test_resident_sdxl.R, and mostly about +# the two places SD 2.1 disagrees with SDXL: the UNet dtype, and a +# component set that names an `encoder` the pipeline never builds. + +library(tinytest) +library(diffuseR) + +# --- dispatch --------------------------------------------------------------------- + +expect_true("sd21" %in% diffuseR:::.resident_families) +expect_true("sd21" %in% eval(formals(resident_load)$model)) +expect_equal(sort(eval(formals(resident_load)$model)), + sort(diffuseR:::.resident_families)) +expect_true(is.function(sd21_load_pipeline)) + +fm <- formals(sd21_load_pipeline) +expect_true(all(c("model_dir", "device", "unet_dtype", "phase_offload", + "verbose") %in% names(fm))) +expect_null(fm$model_dir) +expect_false(fm$phase_offload) +# NULL means the loader decides, and it must decide float32 -- see below. +expect_null(fm$unet_dtype) + +# --- generate-time device injection ------------------------------------------------- + +mk_h <- function(model, device = "cuda:0", gpu = "unet") { + e <- new.env(parent = emptyenv()) + e$model <- model + e$device <- device + e$gpu_components <- gpu + comps <- if (identical(model, "sd21")) { + c("unet", "decoder", "text_encoder") + } else { + c("unet", "decoder", "text_encoder", "text_encoder2") + } + e$staging <- stats::setNames(vector("list", length(comps)), comps) + structure(e, class = "diffuseR_resident") +} + +inj <- diffuseR:::.resident_gen_args(mk_h("sd21"), list()) + +# The device list has to satisfy standardize_devices(), which requires +# every component get_required_components() names -- including `encoder`, +# the VAE encoder img2img needs and the text-to-image pipeline never +# builds. Omitting it is "Missing required component: encoder". +expect_equal(sort(names(inj$devices)), + sort(diffuseR:::get_required_components("sd21"))) +expect_true("encoder" %in% names(inj$devices)) + +# Only the UNet is on the card: SD 2.1's float32 denoise wants ~11.4 GB at +# 768 on its own, so the weights cannot all sit beside it. +expect_equal(inj$devices$unet, "cuda:0") +expect_equal(inj$devices$decoder, "cpu") +expect_equal(inj$devices$text_encoder, "cpu") +# The phantom encoder follows the decoder, which is where it would live. +expect_equal(inj$devices$encoder, "cpu") + +# ... and it keeps following the decoder when the decoder is on the card, +# rather than being pinned to one answer. +wide <- diffuseR:::.resident_gen_args( + mk_h("sd21", gpu = c("unet", "decoder")), list()) +expect_equal(wide$devices$decoder, "cuda:0") +expect_equal(wide$devices$encoder, "cuda:0") +expect_equal(wide$devices$text_encoder, "cpu") + +# The injected list must actually survive the function it is built for. +expect_silent(diffuseR:::standardize_devices( + inj$devices, diffuseR:::get_required_components("sd21"))) + +# An explicit devices= from the caller still wins. +expect_equal(diffuseR:::.resident_gen_args(mk_h("sd21"), + list(devices = list(unet = "cpu")))$devices, + list(unet = "cpu")) + +# SDXL keeps its own four-component shape, no encoder. +sx <- diffuseR:::.resident_gen_args(mk_h("sdxl"), list()) +expect_true("text_encoder2" %in% names(sx$devices)) +expect_equal(sort(names(sx$devices)), + sort(diffuseR:::get_required_components("sdxl"))) + +# --- the float16 trap ---------------------------------------------------------------- + +# SD 2.1's attention overflows in float16 and the pipeline returns all-NaN +# rather than raising, so a float16 resident would load, activate, generate +# and hand back a blank image with nothing to catch. The resident loader +# fixes the dtype before the weights are pinned, so this is the only place +# the decision can be made. +# +# setup_dtype() would answer float16 for a CUDA device, which is why +# sd21_load_pipeline does NOT delegate to it. +expect_equal(diffuseR:::setup_dtype(list(unet = "cuda:0"), NULL), + torch::torch_float16()) + +# Assert the loader's own choice against real weights when they are here. +# Skipped in R CMD check and on a machine without the model. +if (at_home()) { + dir <- file.path(tools::R_user_dir("diffuseR", "data"), "sd21-diffusers") + if (dir.exists(file.path(dir, "unet"))) { + p <- sd21_load_pipeline(model_dir = dir, device = "cuda", + verbose = FALSE) + expect_equal(as.character(p$unet$parameters[[1]]$dtype), "Float") + expect_equal(p$gpu_components, "unet") + expect_false(p$phase_offload) + } +} diff --git a/inst/tinytest/test_resident_sdxl.R b/inst/tinytest/test_resident_sdxl.R index 15945f8..f872c83 100644 --- a/inst/tinytest/test_resident_sdxl.R +++ b/inst/tinytest/test_resident_sdxl.R @@ -120,8 +120,13 @@ expect_equal(diffuseR:::.resident_gpu_set(mk_h("sdxl", gpu = c("unet", "nope"))) inj <- diffuseR:::.resident_gen_args(mk_h("sdxl"), list()) expect_true(!is.null(inj$devices)) +# The list is keyed by what standardize_devices() requires, not by what the +# pipeline builds: SDXL's required set also names an `encoder` (the VAE +# encoder img2img uses) that the text-to-image pipeline never constructs. expect_equal(sort(names(inj$devices)), - c("decoder", "text_encoder", "text_encoder2", "unet")) + sort(diffuseR:::get_required_components("sdxl"))) +expect_true(all(c("unet", "decoder", "text_encoder", "text_encoder2") %in% + names(inj$devices))) # 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 diff --git a/man/resident_load.Rd b/man/resident_load.Rd index daca7fe..8f0bcd1 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -4,7 +4,7 @@ \title{Load a diffusion pipeline as a resident handle} \usage{ resident_load( - model = c("flux2", "flux1", "zimage", "ltx", "sdxl"), + model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), device = "cuda", ..., verbose = TRUE diff --git a/man/sd21_load_pipeline.Rd b/man/sd21_load_pipeline.Rd new file mode 100644 index 0000000..f79c4cb --- /dev/null +++ b/man/sd21_load_pipeline.Rd @@ -0,0 +1,78 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{sd21_load_pipeline} +\alias{sd21_load_pipeline} +\title{Load the SD 2.1 pipeline in the family-loader convention} +\usage{ +sd21_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/}). NULL, the default, resolves the +\code{\link{download_sd21}} 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.} + +\item{unet_dtype}{A torch dtype for the UNet. NULL means float32. Pass +\code{torch::torch_float16()} only if you have reason to believe the +overflow above no longer applies.} + +\item{phase_offload}{Kept for signature parity with the other family +loaders. SD 2.1 has no phased path, so anything but FALSE is ignored.} + +\item{verbose}{Logical.} +} +\value{ +The list from \code{\link{sd_pipeline_from_safetensors}}, plus + \code{phase_offload}. +} +\description{ +The adapter \code{\link{resident_load}} needs, the counterpart to +\code{\link{sdxl_load_pipeline}}. +\code{\link{sd_pipeline_from_safetensors}} takes a required +\code{diffusers_dir} and a plural \code{devices} list; every family +loader takes an optional model directory and a singular \code{device}. +} +\details{ +\strong{The UNet defaults to float32, not float16.} This is the opposite +of \code{\link{sdxl_load_pipeline}} and it is not a preference: SD 2.1's +attention overflows in float16 and the pipeline returns all-NaN, so a +float16 resident would load, activate, generate and hand back a blank +image without raising anything. \code{\link{txt2img_sd21}} already +defaults its native path to float32 for this reason; the resident loader +has to make the same choice, because it fixes the dtype before the +weights are pinned and nothing downstream can undo it. + +Only the UNet goes to the card, as for SDXL, but for a different reason. +SDXL is undone by its float32 VAE decode; SD 2.1 never reaches the +decode. The float32 requirement above makes the denoise itself expensive +-- measured 11.49 GB allocated plus 3.0 GB of allocator slack at the +768x768 default, which does not fit a 15.47 GiB card with the weights +resident as well. That is not a residency artifact: a plain +\code{\link{txt2img_sd21}} with the same devices and dtype OOMs at +11.494 GB, within 0.04 GB of the resident run. + +\code{\link{auto_devices}} already answers "unet on cuda, the rest on +cpu" for this model at every strategy, so residency is agreeing with the +package rather than inventing a placement. + +} +\examples{ +\dontrun{ +res <- resident_load("sd21") +resident_activate(res) +img <- resident_generate(res, "a cat in a spacesuit", seed = 7) +resident_deactivate(res) +} + +} +\seealso{ +\code{\link{resident_load}}, \code{\link{sdxl_load_pipeline}} +}