Add sd21 to the residency contract

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")),

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)

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.

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
 }

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
+}

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,

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)
+    }
+}

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

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