From 152d7c45c6a1ae117d258f5acae6f64bdf2b83d9 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 13:41:33 -0500 Subject: [PATCH 01/18] Add recommend(): unified VRAM + safetensors-read-capability policy recommend(model, vram_gb, st_caps) returns {precision, devices, offload, max_pixels, fork_suggested, note} with per-model tier ladders. nf4 is the default tier; fp8/bf16 upgrade in when the card fits AND the installed safetensors can READ the dtype. .st_can_read is a read probe distinct from the existing write probe (CRAN safetensors reads bf16 it cannot write, so the write probe is the wrong signal for a hosted bf16 artifact); it hand-builds a tiny safetensors file with writeBin and loads it back. When a higher tier fits but is unreadable, recommend() returns nf4 and surfaces the fork suggestion in $note instead of erroring. flux_memory_profile now delegates to recommend("flux1"), fixing the stale tiers that put fp8 (GPU-resident now) in a narrow low-VRAM band it can no longer fit. flux_load_pipeline reconciles the recommendation with the on-disk artifact via .flux_resolve_precision, so a widened fp8 recommendation never points at an unbuilt artifact. --- R/memory_flux.R | 54 +++++---- R/recommend.R | 205 +++++++++++++++++++++++++++++++++ R/st_caps.R | 116 +++++++++++++++++++ R/txt2img_flux.R | 15 ++- inst/tinytest/test_recommend.R | 113 ++++++++++++++++++ 5 files changed, 471 insertions(+), 32 deletions(-) create mode 100644 R/recommend.R create mode 100644 R/st_caps.R create mode 100644 inst/tinytest/test_recommend.R diff --git a/R/memory_flux.R b/R/memory_flux.R index 2b73f62..1b136e9 100644 --- a/R/memory_flux.R +++ b/R/memory_flux.R @@ -1,8 +1,7 @@ #' FLUX Memory Profiles #' -#' VRAM-based execution profiles for the FLUX.1-schnell pipeline, -#' following the LTX-2.3 profile pattern. The 12B transformer runs NF4 -#' (~7 GB, GPU-resident) or fp8 (~12 GB, CPU-resident and streamed); +#' VRAM-based execution profiles for the FLUX.1-schnell pipeline. The +#' 12B transformer runs NF4 (~7 GB) or fp8 (~12 GB), both GPU-resident; #' the T5-XXL text encoder runs float32 on the CPU by default. #' #' @name memory_flux @@ -48,37 +47,36 @@ NULL #' Resolve a FLUX memory profile #' +#' A thin adapter over \code{\link{recommend}} for the FLUX.1 pipeline, +#' kept for back-compatibility. \code{recommend("flux1")} is the policy; +#' this reshapes it into the legacy profile fields the loader consumes. +#' Precision now rises with VRAM (nf4 default, fp8 GPU-resident on 14 GB+ +#' cards when safetensors can read float8, bf16 on 24 GB+); the old +#' bands, which put fp8 in a narrow low-VRAM slot it can no longer fit, +#' were backwards. +#' #' @param vram_gb Numeric or NULL. Available VRAM; auto-detected when #' NULL (via nvidia-smi). #' -#' @return List with \code{name}, \code{precision} ("nf4"/"fp8"), -#' \code{attn_chunk}, \code{text_device}, \code{phase_offload}, and -#' \code{max_pixels} (largest validated image area). +#' @return List with \code{name}, \code{precision} ("nf4"/"fp8"/"bf16"), +#' \code{attn_chunk}, \code{text_device}, \code{phase_offload}, +#' \code{max_pixels}, and (advisory) \code{fork_suggested} and +#' \code{note}. #' #' @export flux_memory_profile <- function(vram_gb = NULL) { - if (is.null(vram_gb)) { - vram_gb <- .detect_vram(use_free = TRUE) - if (is.null(vram_gb) || is.na(vram_gb) || vram_gb <= 0) { - vram_gb <- 0 - } - } - - if (vram_gb >= 12) { - list(name = "high", precision = "nf4", attn_chunk = NULL, - text_device = "cpu", phase_offload = TRUE, - max_pixels = 1536L * 1536L) - } else if (vram_gb >= 9) { - list(name = "medium", precision = "nf4", attn_chunk = 2048L, - text_device = "cpu", phase_offload = TRUE, - max_pixels = 1024L * 1024L) - } else if (vram_gb >= 7) { - list(name = "low", precision = "fp8", attn_chunk = 1024L, - text_device = "cpu", phase_offload = TRUE, - max_pixels = 768L * 768L) + r <- recommend("flux1", vram_gb = vram_gb) + name <- if (identical(r$devices$transformer, "cpu")) { + "cpu_only" + } else if (r$precision %in% c("bf16", "fp8")) { + "high" + } else if (r$max_pixels >= 1024L * 1024L) { + "medium" } else { - list(name = "cpu_only", precision = "nf4", attn_chunk = NULL, - text_device = "cpu", phase_offload = FALSE, - max_pixels = 512L * 512L) + "low" } + list(name = name, precision = r$precision, attn_chunk = r$attn_chunk, + text_device = r$text_device, phase_offload = r$offload, + max_pixels = r$max_pixels, fork_suggested = r$fork_suggested, + note = r$note) } diff --git a/R/recommend.R b/R/recommend.R new file mode 100644 index 0000000..14577cf --- /dev/null +++ b/R/recommend.R @@ -0,0 +1,205 @@ +#' Recommend a precision and device configuration for a model +#' +#' One VRAM-and-capability-aware recommendation for every diffuseR +#' model. The policy: +#' +#' \itemize{ +#' \item nf4 is the default tier. Its weights are packed uint8 plus +#' float32 blocks in sub-2 GB shards, which every safetensors reads, +#' so it always loads. +#' \item When the card has room for a higher-quality tier (fp8 or +#' bf16) AND the installed safetensors can \emph{read} that dtype +#' (\code{\link{.st_can_read}}), that tier is recommended instead. +#' \item When the card has room but safetensors cannot read the tier, +#' nf4 is recommended and the fork suggestion is surfaced in +#' \code{note} (never an error). +#' } +#' +#' This is the policy engine; it does no disk I/O and does not know which +#' artifacts are built. Loaders reconcile the recommendation with what is +#' on disk (see \code{\link{flux_load_pipeline}}). Thresholds are +#' validated on an RTX 5060 Ti (16 GB) and are deliberately conservative +#' elsewhere. Video sizing for \code{"ltx"} is coarse here; the LTX +#' pipeline uses \code{\link{ltx23_memory_profile}} for frame-aware +#' placement. +#' +#' @param model "sd21", "sdxl", "flux1", "flux2", "zimage", or "ltx". +#' @param vram_gb Numeric or NULL. Free VRAM in GB; auto-detected via +#' nvidia-smi when NULL. +#' @param st_caps NULL or a named logical list with \code{bfloat16} +#' and/or \code{float8_e4m3fn} - the safetensors READ capabilities. +#' NULL probes the installed safetensors. +#' +#' @return A list with \code{model}, \code{precision}, \code{devices} +#' (named component -> device map), \code{offload} (phase-offloading +#' logical), \code{max_pixels}, \code{text_device}, \code{attn_chunk}, +#' \code{vram_gb}, \code{fork_suggested} (logical), and \code{note} +#' (the fork suggestion string, or NULL). +#' +#' @export +#' +#' @examples +#' \dontrun{ +#' # Auto-detect VRAM and probe the installed safetensors +#' recommend("flux2") +#' +#' # A 16 GB card without float8 support: fp8 wanted, nf4 recommended +#' r <- recommend("flux1", vram_gb = 16, +#' st_caps = list(bfloat16 = TRUE, float8_e4m3fn = FALSE)) +#' r$precision # "nf4" +#' r$fork_suggested # TRUE +#' cat(r$note) # the fork-or-nf4 message +#' } +recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", + "ltx"), + vram_gb = NULL, st_caps = NULL) { + model <- match.arg(model) + if (is.null(vram_gb)) { + vram_gb <- .detect_vram(use_free = TRUE) + } + if (is.null(vram_gb) || is.na(vram_gb) || vram_gb < 0) { + vram_gb <- 0 + } + if (is.null(st_caps)) { + st_caps <- list(bfloat16 = .st_can_read("bfloat16"), + float8_e4m3fn = .st_can_read("float8_e4m3fn")) + } + + tiers <- .recommend_specs()[[model]] + + chosen <- NULL + want <- NULL # first VRAM-eligible tier blocked by a missing read cap + for (tier in tiers) { + if (vram_gb < tier$min_vram) { + next + } + need <- tier$needs + if (is.null(need) || isTRUE(st_caps[[need]])) { + chosen <- tier + break + } + if (is.null(want)) { + want <- tier + } + } + if (is.null(chosen)) { + chosen <- tiers[[length(tiers)]] # terminal cpu tier + } + + fork <- !is.null(want) && !identical(want$precision, chosen$precision) + list( + model = model, + precision = chosen$precision, + devices = chosen$devices, + offload = isTRUE(chosen$offload), + max_pixels = chosen$max_pixels, + text_device = chosen$text_device %||% "cpu", + attn_chunk = chosen$attn_chunk, + vram_gb = vram_gb, + fork_suggested = fork, + note = if (fork) .st_fork_note(want$precision) else NULL + ) +} + +# flux-family component placement: the big DiT and the VAE compute on +# the GPU (or all on CPU for the cpu tier); the text encoder is resident +# on the CPU and phase-onloaded during its own phase. +.dev_flux <- function(gpu = TRUE) { + if (gpu) { + list(transformer = "cuda", text = "cpu", vae = "cuda") + } else { + list(transformer = "cpu", text = "cpu", vae = "cpu") + } +} + +# One bf16/fp8/nf4/nf4-tight/cpu ladder for the flux-family image models. +# Precision rises with VRAM (nf4 default, fp8/bf16 as upgrades) - the +# inverse of the old flux_memory_profile, which had fp8 in a narrow +# low-VRAM band it can no longer fit now that fp8 is GPU-resident. +.flux_family_tiers <- function(bf16_vram, fp8_vram, nf4_vram, nf4_tight_vram, + max_hi, max_mid, max_lo, max_cpu, + attn_tight = NULL) { + list( + list(precision = "bf16", min_vram = bf16_vram, needs = "bfloat16", + devices = .dev_flux(TRUE), offload = TRUE, max_pixels = max_hi, + attn_chunk = NULL), + list(precision = "fp8", min_vram = fp8_vram, needs = "float8_e4m3fn", + devices = .dev_flux(TRUE), offload = TRUE, max_pixels = max_mid, + attn_chunk = NULL), + list(precision = "nf4", min_vram = nf4_vram, needs = NULL, + devices = .dev_flux(TRUE), offload = TRUE, max_pixels = max_mid, + attn_chunk = NULL), + list(precision = "nf4", min_vram = nf4_tight_vram, needs = NULL, + devices = .dev_flux(TRUE), offload = TRUE, max_pixels = max_lo, + attn_chunk = attn_tight), + list(precision = "nf4", min_vram = 0, needs = NULL, cpu = TRUE, + devices = .dev_flux(FALSE), offload = FALSE, max_pixels = max_cpu, + attn_chunk = NULL) + ) +} + +# SD-family ladder: fp16 for cards that fit the full model, nf4 default +# for tighter cards, cpu otherwise. Both dtypes are CRAN-readable, so no +# fork gate. Device maps reuse the auto_devices strategy builder. +.sd_tiers <- function(model, fp16_vram, nf4_vram, max_fp16, max_nf4, max_cpu) { + list( + list(precision = "fp16", min_vram = fp16_vram, needs = NULL, + devices = .build_fallback_devices(model, "full_gpu"), + offload = FALSE, max_pixels = max_fp16), + list(precision = "nf4", min_vram = nf4_vram, needs = NULL, + devices = .build_fallback_devices(model, "unet_gpu"), + offload = TRUE, max_pixels = max_nf4), + list(precision = "nf4", min_vram = 0, needs = NULL, cpu = TRUE, + devices = .build_fallback_devices(model, "cpu_only"), + offload = FALSE, max_pixels = max_cpu) + ) +} + +# Per-model tier ladders. Built lazily (device maps call helpers) so the +# thresholds live in one auditable place. +.recommend_specs <- function() { + px <- function(n) as.integer(n) * as.integer(n) + list( + sd21 = .sd_tiers("sd21", fp16_vram = 6, nf4_vram = 3, + max_fp16 = px(1024), max_nf4 = px(768), + max_cpu = px(512)), + sdxl = .sd_tiers("sdxl", fp16_vram = 12, nf4_vram = 6, + max_fp16 = px(1024), max_nf4 = px(1024), + max_cpu = px(768)), + # 12B: nf4 peaks ~9.6 GB at 1024^2, fp8 ~12 GB resident, + # bf16 needs a 24 GB card. attn-chunk the tight nf4 tier. + flux1 = .flux_family_tiers(bf16_vram = 24, fp8_vram = 14, + nf4_vram = 10, nf4_tight_vram = 8, + max_hi = px(1536), max_mid = px(1024), + max_lo = px(768), max_cpu = px(512), + attn_tight = 2048L), + # 4B but activation-heavy: 1024^2 peaks ~12.5 GB regardless of + # weight precision, so the 1024^2 tiers want ~13 GB free. + flux2 = .flux_family_tiers(bf16_vram = 16, fp8_vram = 14, + nf4_vram = 13, nf4_tight_vram = 8, + max_hi = px(1024), max_mid = px(1024), + max_lo = px(768), max_cpu = px(512)), + # 6B: 1024^2 peaks ~13.1 GB, 512^2 ~half that. + zimage = .flux_family_tiers(bf16_vram = 18, fp8_vram = 14, + nf4_vram = 13, nf4_tight_vram = 8, + max_hi = px(1024), max_mid = px(1024), + max_lo = px(512), max_cpu = px(512)), + # 22B video. fp8 is CPU-resident and streamed here (unlike the + # image models); nf4 keeps the transformer resident. Coarse - + # ltx23_memory_profile does the frame-aware sizing. + ltx = list( + list(precision = "nf4", min_vram = 14, needs = NULL, + devices = .dev_flux(TRUE), offload = TRUE, + max_pixels = px(1280), text_device = "cpu", + attn_chunk = NULL), + list(precision = "fp8", min_vram = 10, + needs = "float8_e4m3fn", devices = .dev_flux(TRUE), + offload = TRUE, max_pixels = px(1024), + text_device = "cpu", attn_chunk = 4096L), + list(precision = "nf4", min_vram = 0, needs = NULL, + cpu = TRUE, devices = .dev_flux(FALSE), + offload = FALSE, max_pixels = px(512), + text_device = "cpu", attn_chunk = NULL) + ) + ) +} diff --git a/R/st_caps.R b/R/st_caps.R new file mode 100644 index 0000000..3732fec --- /dev/null +++ b/R/st_caps.R @@ -0,0 +1,116 @@ +#' safetensors read-capability probes and fork messaging +#' +#' CRAN safetensors (<= 0.2.1) reads bfloat16 but cannot write it, and +#' has no float8 support at all; the fixes are upstream +#' (mlverse/safetensors#11 for bfloat16 write, #13 for float8) and in the +#' cornball-ai/safetensors fork. Two capabilities matter and they differ: +#' +#' \itemize{ +#' \item \emph{write} (\code{\link{flux_quantize}}'s internal +#' \code{.st_can_write}, in quantize_flux.R): needed to BUILD a +#' quantized artifact in that dtype. +#' \item \emph{read} (\code{.st_can_read}, here): needed to LOAD a +#' hosted artifact in that dtype. This is the capability that gates +#' user-facing recommendations. It is strictly weaker than write: +#' CRAN safetensors reads bfloat16 it cannot write, so the write +#' probe is the wrong signal for whether a hosted bf16 artifact will +#' load. +#' } +#' +#' Both are capability-probed, never version-pinned, so the fork +#' requirement self-heals the day the fixes reach CRAN. +#' +#' @name st_caps +NULL + +# Read-probe cache, keyed by dtype. Separate from quantize_flux.R's +# `.st_caps` write cache: the same dtype can be readable but not +# writable (bfloat16 on CRAN), so the two must not share entries. +.st_read_caps <- new.env(parent = emptyenv()) + +# Write a minimal 2-element safetensors file by hand: a u64 +# little-endian header length, the JSON header, then the raw tensor +# bytes. Deliberately does NOT go through safetensors::safe_save_file - +# that is the whole point, since a CRAN safetensors cannot WRITE +# bfloat16 yet can READ it. Lets `.st_can_read` test read capability in +# isolation from write capability. +.st_write_min <- function(path, dtype_name, payload) { + hdr <- sprintf('{"w":{"dtype":"%s","shape":[2],"data_offsets":[0,%d]}}', + dtype_name, length(payload)) + hb <- charToRaw(hdr) + n <- length(hb) + # u64 header length, little-endian (headers here are < 256 bytes) + len_bytes <- as.raw(c(n %% 256L, (n %/% 256L) %% 256L, 0L, 0L, 0L, 0L, + 0L, 0L)) + con <- file(path, "wb") + on.exit(close(con), add = TRUE) + writeBin(len_bytes, con) + writeBin(hb, con) + writeBin(payload, con) + invisible(path) +} + +# The [0, 1] byte patterns for the two dtypes the recommender gates on. +# bfloat16 is the top 16 bits of float32: 1.0f = 0x3F800000 -> 0x3F80 +# (little-endian 0x80 0x3F); 0.0 -> 0x0000. float8_e4m3fn: 1.0 = exp +# bias 7, mantissa 0 = 0x38; 0.0 = 0x00. +.st_read_probe_spec <- list( + bfloat16 = list(name = "BF16", + bytes = as.raw(c(0x00, 0x00, 0x80, 0x3f))), + float8_e4m3fn = list(name = "F8_E4M3", + bytes = as.raw(c(0x00, 0x38))) +) + +#' Probe whether the installed safetensors can READ a dtype +#' +#' Hand-builds a tiny safetensors file of the dtype (via +#' \code{.st_write_min}, no safetensors writer involved) and tries to +#' load it back. Cached per session; +#' \code{options(diffuseR.st_read_caps = list(bfloat16 = TRUE, ...))} +#' overrides the probe for tests and for forcing a tier. +#' +#' @param dtype "bfloat16" or "float8_e4m3fn". +#' @return Logical. +#' @keywords internal +.st_can_read <- function(dtype = c("bfloat16", "float8_e4m3fn")) { + dtype <- match.arg(dtype) + override <- getOption("diffuseR.st_read_caps") + if (!is.null(override) && !is.null(override[[dtype]])) { + return(isTRUE(override[[dtype]])) + } + cached <- .st_read_caps[[dtype]] + if (!is.null(cached)) { + return(cached) + } + ok <- requireNamespace("safetensors", quietly = TRUE) && tryCatch({ + spec <- .st_read_probe_spec[[dtype]] + tmp <- tempfile(fileext = ".safetensors") + on.exit(unlink(tmp), add = TRUE) + .st_write_min(tmp, spec$name, spec$bytes) + y <- safetensors::safe_load_file(tmp, framework = "torch") + !is.null(y$w) && identical(as.integer(y$w$shape), 2L) + }, error = function(e) FALSE) + .st_read_caps[[dtype]] <- ok + ok +} + +# The standard "install the fork, or press on with nf4" message. Shared +# by the recommender (read side) and the download graceful-fallback path +# (write side) so the wording stays identical everywhere. No em dashes +# (house style). +.st_fork_note <- function(precision) { + precision <- as.character(precision) + detail <- switch(precision, + fp8 = "float8 support (mlverse/safetensors#13)", + float8_e4m3fn = "float8 support (mlverse/safetensors#13)", + bf16 = "bfloat16 write support (mlverse/safetensors#11)", + bfloat16 = "bfloat16 write support (mlverse/safetensors#11)", + paste0(precision, " support (mlverse/safetensors)")) + sprintf(paste0( + "%s is the best fit for your card but needs ", + "cornball-ai/safetensors until CRAN safetensors ships %s. ", + "Install remotes::install_github(\"cornball-ai/safetensors\"), ", + "or press on with nf4: same weights, slightly lower ", + "precision, and it just works."), + precision, detail) +} diff --git a/R/txt2img_flux.R b/R/txt2img_flux.R index 857489f..094855a 100644 --- a/R/txt2img_flux.R +++ b/R/txt2img_flux.R @@ -98,15 +98,22 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", attn_chunk = NULL, phase_offload = TRUE, verbose = TRUE) { profile <- flux_memory_profile() - if (is.null(precision)) { - precision <- profile$precision + if (verbose && !is.null(profile$note)) { + message(profile$note) } if (is.null(attn_chunk)) { attn_chunk <- profile$attn_chunk } + prefix <- file.path(tools::R_user_dir("diffuseR", "data"), + "flux1-schnell-") + if (is.null(precision)) { + # Reconcile the VRAM recommendation with what is on disk: prefer + # a built artifact (fp8 first when readable), else nf4. Keeps a + # widened fp8 recommendation from pointing at an unbuilt artifact. + precision <- .flux_resolve_precision("auto", prefix) + } if (is.null(model_dir)) { - model_dir <- file.path(tools::R_user_dir("diffuseR", "data"), - paste0("flux1-schnell-", precision)) + model_dir <- paste0(prefix, precision) } ckpt <- if (file.exists(file.path(model_dir, "manifest.json"))) { diff --git a/inst/tinytest/test_recommend.R b/inst/tinytest/test_recommend.R new file mode 100644 index 0000000..7611093 --- /dev/null +++ b/inst/tinytest/test_recommend.R @@ -0,0 +1,113 @@ +# recommend(): VRAM + safetensors-read-capability policy, the read +# probe, and the flux_memory_profile adapter. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +library(diffuseR) + +fork <- list(bfloat16 = TRUE, float8_e4m3fn = TRUE) # can read everything +cran <- list(bfloat16 = TRUE, float8_e4m3fn = FALSE) # reads bf16, not fp8 + +# --- structure -------------------------------------------------------------------- + +r <- recommend("flux1", vram_gb = 16, st_caps = fork) +expect_true(is.list(r)) +expect_equal(r$model, "flux1") +expect_true(all(c("precision", "devices", "offload", "max_pixels", + "text_device", "attn_chunk", "vram_gb", "fork_suggested", "note") %in% + names(r))) +expect_true(is.integer(r$max_pixels) || is.numeric(r$max_pixels)) + +# --- nf4 is the floor / default --------------------------------------------------- + +# No GPU: nf4 on CPU, no fork nag, smallest area. +for (m in c("sd21", "sdxl", "flux1", "flux2", "zimage")) { + z <- recommend(m, vram_gb = 0, st_caps = cran) + expect_equal(z$precision, "nf4") + expect_false(z$fork_suggested) + expect_true(all(unlist(z$devices) == "cpu")) +} + +# Mid VRAM with no float8: nf4, and no nag when the card could not fit +# fp8 anyway (fp8 tier not VRAM-eligible). +expect_equal(recommend("flux1", 12, cran)$precision, "nf4") +expect_false(recommend("flux1", 12, cran)$fork_suggested) + +# --- fp8 upgrade + fork gating ---------------------------------------------------- + +# 16 GB with float8 read: fp8 recommended, no nag. +f16 <- recommend("flux1", 16, fork) +expect_equal(f16$precision, "fp8") +expect_false(f16$fork_suggested) +expect_null(f16$note) +expect_equal(f16$devices$transformer, "cuda") + +# 16 GB WITHOUT float8 read: fp8 wanted but unreadable -> nf4 + fork nag. +c16 <- recommend("flux1", 16, cran) +expect_equal(c16$precision, "nf4") +expect_true(c16$fork_suggested) +expect_true(is.character(c16$note) && grepl("safetensors#13", c16$note)) +expect_false(grepl("—", c16$note)) # no em dash (house style) + +# --- bf16 top tier (CRAN-readable, no fork needed) -------------------------------- + +# 24 GB reads bf16 on CRAN too -> bf16, no nag. +b24 <- recommend("flux1", 24, cran) +expect_equal(b24$precision, "bf16") +expect_false(b24$fork_suggested) + +# But if bf16 read were somehow missing, it is gated like fp8. +nob <- recommend("flux1", 24, list(bfloat16 = FALSE, float8_e4m3fn = FALSE)) +expect_true(nob$fork_suggested) +expect_equal(nob$precision, "nf4") + +# --- SD ladder: fp16 for cards that fit, nf4 default, cpu ------------------------- + +expect_equal(recommend("sdxl", 16, cran)$precision, "fp16") +expect_equal(recommend("sdxl", 8, cran)$precision, "nf4") +expect_equal(recommend("sdxl", 16, cran)$devices$unet, "cuda") +expect_true("text_encoder2" %in% names(recommend("sdxl", 16, cran)$devices)) +expect_false("text_encoder2" %in% names(recommend("sd21", 16, cran)$devices)) +# SD is CRAN-readable at every tier: never a fork nag. +for (v in c(0, 4, 8, 16, 24)) expect_false(recommend("sdxl", v, cran)$fork_suggested) + +# --- precision rises monotonically with VRAM (quality never drops) ---------------- + +rank <- c(nf4 = 1L, fp16 = 2L, fp8 = 3L, bf16 = 4L) +for (m in c("flux1", "flux2", "zimage")) { + precs <- vapply(c(0, 8, 12, 16, 24), + function(v) rank[[recommend(m, v, fork)$precision]], integer(1)) + expect_true(!is.unsorted(precs)) +} + +# --- read probe: override option wins, and self-consistency ----------------------- + +expect_true(is.logical(diffuseR:::.st_can_read("bfloat16"))) +options(diffuseR.st_read_caps = list(bfloat16 = FALSE, float8_e4m3fn = FALSE)) +expect_false(diffuseR:::.st_can_read("bfloat16")) +expect_false(diffuseR:::.st_can_read("float8_e4m3fn")) +# recommend() with st_caps = NULL now reads the (forced) probe. +expect_true(recommend("flux1", 24)$fork_suggested) +options(diffuseR.st_read_caps = NULL) + +# On a machine whose safetensors CAN read a dtype, the hand-built probe +# file must decode (guards the writeBin byte patterns). +if (requireNamespace("safetensors", quietly = TRUE)) { + # These reflect the installed safetensors; just assert they are logical + # and, when TRUE, that recommend trusts them. + expect_true(is.logical(diffuseR:::.st_can_read("float8_e4m3fn"))) +} + +# --- flux_memory_profile adapter follows recommend -------------------------------- + +options(diffuseR.st_read_caps = fork) +p <- flux_memory_profile(vram_gb = 16) +expect_equal(p$precision, "fp8") +expect_equal(p$phase_offload, TRUE) +expect_true(p$max_pixels >= 1024L * 1024L) +options(diffuseR.st_read_caps = cran) +p2 <- flux_memory_profile(vram_gb = 16) +expect_equal(p2$precision, "nf4") +expect_true(isTRUE(p2$fork_suggested)) +options(diffuseR.st_read_caps = NULL) From 73959825dbb0fb90176bf57d5727ba09f744abcb Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 13:44:11 -0500 Subject: [PATCH 02/18] Default quantizer shards to <2 GB (CRAN-safetensors readable) R safetensors overflows a 32-bit offset on files at/above 2^31 bytes, so the old 4e9 shard default produced shards stock CRAN safetensors cannot read back (only the cornball-ai fork could). flux_quantize and the LTX nf4/fp8 writers now default shard_bytes to 1.9e9, keeping every shard under the ~2.15 GB ceiling with headroom for the post-add flush overshoot. 4e9 stays available via the argument for local fork builds. --- R/fp8_ltx23.R | 7 +++++-- R/nf4_ltx23.R | 7 +++++-- R/quantize_flux.R | 8 ++++++-- 3 files changed, 16 insertions(+), 6 deletions(-) diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index 88233ba..a6e737d 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -105,7 +105,10 @@ ltx23_fp8_linear <- torch::nn_module( #' #' @param checkpoint_path Source .safetensors (46 GB bf16 single file). #' @param output_dir Output directory for shards + manifest. -#' @param shard_bytes Numeric. Approximate shard size (default 4 GB). +#' @param shard_bytes Numeric. Target shard size in bytes. The default +#' 1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +#' stock CRAN safetensors can read. Pass a larger value (e.g. 4e9) only +#' for local builds you will read back with a fork-patched safetensors. #' @param force Logical. Re-quantize even if a valid manifest exists. #' @param verbose Logical. #' @@ -114,7 +117,7 @@ ltx23_fp8_linear <- torch::nn_module( #' @export ltx23_quantize_fp8 <- function(checkpoint_path, output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), - shard_bytes = 4e9, force = FALSE, + shard_bytes = 1.9e9, force = FALSE, verbose = TRUE) { manifest_path <- file.path(output_dir, "manifest.json") if (!force && file.exists(manifest_path)) { diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R index 4c454a9..9166ccb 100644 --- a/R/nf4_ltx23.R +++ b/R/nf4_ltx23.R @@ -273,7 +273,10 @@ ltx23_nf4_linear <- torch::nn_module( #' #' @param checkpoint_path Source .safetensors (bf16 single file). #' @param output_dir Output directory for shards + manifest. -#' @param shard_bytes Numeric. Approximate shard size. +#' @param shard_bytes Numeric. Target shard size in bytes. The default +#' 1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +#' stock CRAN safetensors can read. Pass a larger value (e.g. 4e9) only +#' for local builds you will read back with a fork-patched safetensors. #' @param force Logical. Re-quantize even if a valid manifest exists. #' @param verbose Logical. #' @@ -282,7 +285,7 @@ ltx23_nf4_linear <- torch::nn_module( #' @export ltx23_quantize_nf4 <- function(checkpoint_path, output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-nf4"), - shard_bytes = 4e9, force = FALSE, + shard_bytes = 1.9e9, force = FALSE, verbose = TRUE) { manifest_path <- file.path(output_dir, "manifest.json") if (!force && file.exists(manifest_path)) { diff --git a/R/quantize_flux.R b/R/quantize_flux.R index c47cc9d..caa5238 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -200,7 +200,11 @@ NULL #' @param output_dir Output directory for shards + manifest (default: #' the per-format location under \code{tools::R_user_dir}). #' @param format "nf4" or "fp8". -#' @param shard_bytes Numeric. Approximate shard size. +#' @param shard_bytes Numeric. Target shard size in bytes. The default +#' 1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +#' stock CRAN safetensors can read, so the artifact loads fork-free. +#' Pass a larger value (e.g. 4e9) only for local builds you will read +#' back with a fork-patched safetensors. #' @param force Logical. Re-quantize even if a valid manifest exists. #' @param verbose Logical. #' @@ -208,7 +212,7 @@ NULL #' #' @export flux_quantize <- function(transformer_dir, output_dir = NULL, - format = c("nf4", "fp8"), shard_bytes = 4e9, + format = c("nf4", "fp8"), shard_bytes = 1.9e9, force = FALSE, verbose = TRUE) { format <- match.arg(format) if (format == "fp8" && !.st_can_write("float8_e4m3fn")) { From b3a50c6b436ad641e3e5248364dfc4306d8ee746 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 14:16:37 -0500 Subject: [PATCH 03/18] document(): regenerate man + NAMESPACE for recommend/st_caps and shard_bytes docs --- NAMESPACE | 1 + man/dot-st_can_read.Rd | 21 +++++++++++++ man/flux_memory_profile.Rd | 15 ++++++--- man/flux_quantize.Rd | 8 +++-- man/ltx23_quantize_fp8.Rd | 7 +++-- man/ltx23_quantize_nf4.Rd | 7 +++-- man/memory_flux.Rd | 5 ++- man/recommend.Rd | 64 ++++++++++++++++++++++++++++++++++++++ man/st_caps.Rd | 27 ++++++++++++++++ 9 files changed, 142 insertions(+), 13 deletions(-) create mode 100644 man/dot-st_can_read.Rd create mode 100644 man/recommend.Rd create mode 100644 man/st_caps.Rd diff --git a/NAMESPACE b/NAMESPACE index 1d44233..ae1eda3 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -169,6 +169,7 @@ export(preprocess_image) export(quant_conv) export(qwen_bpe_tokenizer) export(qwen3_encoder) +export(recommend) export(save_frames) export(save_image) export(save_video) diff --git a/man/dot-st_can_read.Rd b/man/dot-st_can_read.Rd new file mode 100644 index 0000000..0381b9c --- /dev/null +++ b/man/dot-st_can_read.Rd @@ -0,0 +1,21 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{.st_can_read} +\alias{.st_can_read} +\title{Probe whether the installed safetensors can READ a dtype} +\usage{ +.st_can_read(dtype = c("bfloat16", "float8_e4m3fn")) +} +\arguments{ +\item{dtype}{"bfloat16" or "float8_e4m3fn".} +} +\value{ +Logical. +} +\description{ +Hand-builds a tiny safetensors file of the dtype (via +\code{.st_write_min}, no safetensors writer involved) and tries to +load it back. Cached per session; +\code{options(diffuseR.st_read_caps = list(bfloat16 = TRUE, ...))} +overrides the probe for tests and for forcing a tier. +} +\keyword{internal} diff --git a/man/flux_memory_profile.Rd b/man/flux_memory_profile.Rd index a9474fe..88f55b4 100644 --- a/man/flux_memory_profile.Rd +++ b/man/flux_memory_profile.Rd @@ -10,10 +10,17 @@ flux_memory_profile(vram_gb = NULL) NULL (via nvidia-smi).} } \value{ -List with \code{name}, \code{precision} ("nf4"/"fp8"), - \code{attn_chunk}, \code{text_device}, \code{phase_offload}, and - \code{max_pixels} (largest validated image area). +List with \code{name}, \code{precision} ("nf4"/"fp8"/"bf16"), + \code{attn_chunk}, \code{text_device}, \code{phase_offload}, + \code{max_pixels}, and (advisory) \code{fork_suggested} and + \code{note}. } \description{ -Resolve a FLUX memory profile +A thin adapter over \code{\link{recommend}} for the FLUX.1 pipeline, +kept for back-compatibility. \code{recommend("flux1")} is the policy; +this reshapes it into the legacy profile fields the loader consumes. +Precision now rises with VRAM (nf4 default, fp8 GPU-resident on 14 GB+ +cards when safetensors can read float8, bf16 on 24 GB+); the old +bands, which put fp8 in a narrow low-VRAM slot it can no longer fit, +were backwards. } diff --git a/man/flux_quantize.Rd b/man/flux_quantize.Rd index 4f97593..bcac1b1 100644 --- a/man/flux_quantize.Rd +++ b/man/flux_quantize.Rd @@ -4,7 +4,7 @@ \title{Quantize a FLUX transformer to NF4 or fp8 shards} \usage{ flux_quantize(transformer_dir, output_dir = NULL, format = c("nf4", "fp8"), - shard_bytes = 4e+09, force = FALSE, verbose = TRUE) + shard_bytes = 1.9e+09, force = FALSE, verbose = TRUE) } \arguments{ \item{transformer_dir}{Source diffusers transformer directory.} @@ -14,7 +14,11 @@ the per-format location under \code{tools::R_user_dir}).} \item{format}{"nf4" or "fp8".} -\item{shard_bytes}{Numeric. Approximate shard size.} +\item{shard_bytes}{Numeric. Target shard size in bytes. The default +1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +stock CRAN safetensors can read, so the artifact loads fork-free. +Pass a larger value (e.g. 4e9) only for local builds you will read +back with a fork-patched safetensors.} \item{force}{Logical. Re-quantize even if a valid manifest exists.} diff --git a/man/ltx23_quantize_fp8.Rd b/man/ltx23_quantize_fp8.Rd index 3ed610c..69d6bc1 100644 --- a/man/ltx23_quantize_fp8.Rd +++ b/man/ltx23_quantize_fp8.Rd @@ -5,14 +5,17 @@ \usage{ ltx23_quantize_fp8(checkpoint_path, output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-fp8"), - shard_bytes = 4e+09, force = FALSE, verbose = TRUE) + shard_bytes = 1.9e+09, force = FALSE, verbose = TRUE) } \arguments{ \item{checkpoint_path}{Source .safetensors (46 GB bf16 single file).} \item{output_dir}{Output directory for shards + manifest.} -\item{shard_bytes}{Numeric. Approximate shard size (default 4 GB).} +\item{shard_bytes}{Numeric. Target shard size in bytes. The default +1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +stock CRAN safetensors can read. Pass a larger value (e.g. 4e9) only +for local builds you will read back with a fork-patched safetensors.} \item{force}{Logical. Re-quantize even if a valid manifest exists.} diff --git a/man/ltx23_quantize_nf4.Rd b/man/ltx23_quantize_nf4.Rd index 14c508b..4d72932 100644 --- a/man/ltx23_quantize_nf4.Rd +++ b/man/ltx23_quantize_nf4.Rd @@ -5,14 +5,17 @@ \usage{ ltx23_quantize_nf4(checkpoint_path, output_dir = file.path(tools::R_user_dir("diffuseR", "data"), "ltx2.3-nf4"), - shard_bytes = 4e+09, force = FALSE, verbose = TRUE) + shard_bytes = 1.9e+09, force = FALSE, verbose = TRUE) } \arguments{ \item{checkpoint_path}{Source .safetensors (bf16 single file).} \item{output_dir}{Output directory for shards + manifest.} -\item{shard_bytes}{Numeric. Approximate shard size.} +\item{shard_bytes}{Numeric. Target shard size in bytes. The default +1.9e9 keeps every shard under the 2^31-byte (~2.15 GB) ceiling that +stock CRAN safetensors can read. Pass a larger value (e.g. 4e9) only +for local builds you will read back with a fork-patched safetensors.} \item{force}{Logical. Re-quantize even if a valid manifest exists.} diff --git a/man/memory_flux.Rd b/man/memory_flux.Rd index 4c79371..083121d 100644 --- a/man/memory_flux.Rd +++ b/man/memory_flux.Rd @@ -3,8 +3,7 @@ \alias{memory_flux} \title{FLUX Memory Profiles} \description{ -VRAM-based execution profiles for the FLUX.1-schnell pipeline, -following the LTX-2.3 profile pattern. The 12B transformer runs NF4 -(~7 GB, GPU-resident) or fp8 (~12 GB, CPU-resident and streamed); +VRAM-based execution profiles for the FLUX.1-schnell pipeline. The +12B transformer runs NF4 (~7 GB) or fp8 (~12 GB), both GPU-resident; the T5-XXL text encoder runs float32 on the CPU by default. } diff --git a/man/recommend.Rd b/man/recommend.Rd new file mode 100644 index 0000000..79ed32e --- /dev/null +++ b/man/recommend.Rd @@ -0,0 +1,64 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{recommend} +\alias{recommend} +\title{Recommend a precision and device configuration for a model} +\usage{ +recommend(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", "ltx"), + vram_gb = NULL, st_caps = NULL) +} +\arguments{ +\item{model}{"sd21", "sdxl", "flux1", "flux2", "zimage", or "ltx".} + +\item{vram_gb}{Numeric or NULL. Free VRAM in GB; auto-detected via +nvidia-smi when NULL.} + +\item{st_caps}{NULL or a named logical list with \code{bfloat16} +and/or \code{float8_e4m3fn} - the safetensors READ capabilities. +NULL probes the installed safetensors.} +} +\value{ +A list with \code{model}, \code{precision}, \code{devices} + (named component -> device map), \code{offload} (phase-offloading + logical), \code{max_pixels}, \code{text_device}, \code{attn_chunk}, + \code{vram_gb}, \code{fork_suggested} (logical), and \code{note} + (the fork suggestion string, or NULL). +} +\description{ +One VRAM-and-capability-aware recommendation for every diffuseR +model. The policy: +} +\details{ +\itemize{ +\item nf4 is the default tier. Its weights are packed uint8 plus +float32 blocks in sub-2 GB shards, which every safetensors reads, +so it always loads. +\item When the card has room for a higher-quality tier (fp8 or +bf16) AND the installed safetensors can \emph{read} that dtype +(\code{\link{.st_can_read}}), that tier is recommended instead. +\item When the card has room but safetensors cannot read the tier, +nf4 is recommended and the fork suggestion is surfaced in +\code{note} (never an error). +} + +This is the policy engine; it does no disk I/O and does not know which +artifacts are built. Loaders reconcile the recommendation with what is +on disk (see \code{\link{flux_load_pipeline}}). Thresholds are +validated on an RTX 5060 Ti (16 GB) and are deliberately conservative +elsewhere. Video sizing for \code{"ltx"} is coarse here; the LTX +pipeline uses \code{\link{ltx23_memory_profile}} for frame-aware +placement. + +} +\examples{ +\dontrun{ +# Auto-detect VRAM and probe the installed safetensors +recommend("flux2") + +# A 16 GB card without float8 support: fp8 wanted, nf4 recommended +r <- recommend("flux1", vram_gb = 16, + st_caps = list(bfloat16 = TRUE, float8_e4m3fn = FALSE)) +r$precision # "nf4" +r$fork_suggested # TRUE +cat(r$note) # the fork-or-nf4 message +} +} diff --git a/man/st_caps.Rd b/man/st_caps.Rd new file mode 100644 index 0000000..2dfbb1d --- /dev/null +++ b/man/st_caps.Rd @@ -0,0 +1,27 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{st_caps} +\alias{st_caps} +\title{safetensors read-capability probes and fork messaging} +\description{ +CRAN safetensors (<= 0.2.1) reads bfloat16 but cannot write it, and +has no float8 support at all; the fixes are upstream +(mlverse/safetensors#11 for bfloat16 write, #13 for float8) and in the +cornball-ai/safetensors fork. Two capabilities matter and they differ: +} +\details{ +\itemize{ +\item \emph{write} (\code{\link{flux_quantize}}'s internal +\code{.st_can_write}, in quantize_flux.R): needed to BUILD a +quantized artifact in that dtype. +\item \emph{read} (\code{.st_can_read}, here): needed to LOAD a +hosted artifact in that dtype. This is the capability that gates +user-facing recommendations. It is strictly weaker than write: +CRAN safetensors reads bfloat16 it cannot write, so the write +probe is the wrong signal for whether a hosted bf16 artifact will +load. +} + +Both are capability-probed, never version-pinned, so the fork +requirement self-heals the day the fixes reach CRAN. + +} From 9f24697a3c6cdb00395224f9f4e9e9871b6c87f6 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 14:45:32 -0500 Subject: [PATCH 04/18] Graceful precision fallback + multi-GB read breadcrumb When a user explicitly requests fp8/bf16 and the safetensors capability is missing, .st_graceful_precision prints the fork suggestion and falls back to nf4 instead of erroring; wired into download_flux1/flux2/zimage. The low-level flux_quantize/flux_load_transformer keep their hard errors (builder primitives, and a test pins that). Checkpoint reads route through .st_read_or_breadcrumb: a read that fails on a shard at/above 2^31 bytes is retranslated from the cryptic 32-bit overflow into an actionable rebuild-with-smaller-shards-or-install-fork message. Reactive (only on a real failure of an oversize shard), so it never false-alarms on the fork or on sub-2 GB shards. --- R/checkpoint_flux.R | 7 ++- R/download_flux.R | 3 ++ R/download_flux2.R | 2 + R/download_zimage.R | 2 + R/fp8_ltx23.R | 11 +++-- R/st_caps.R | 87 ++++++++++++++++++++++++++++++---- inst/tinytest/test_recommend.R | 45 ++++++++++++++++++ 7 files changed, 142 insertions(+), 15 deletions(-) diff --git a/R/checkpoint_flux.R b/R/checkpoint_flux.R index fc9f551..6f99524 100644 --- a/R/checkpoint_flux.R +++ b/R/checkpoint_flux.R @@ -170,7 +170,8 @@ flux_open_checkpoint <- function(transformer_dir) { if (is.null(shard) || is.na(shard)) { stop("Key not found in checkpoint index: ", key) } - handles[[shard]]$get_tensor(key) + .st_read_or_breadcrumb(function() handles[[shard]]$get_tensor(key), + file.path(dir, shard)) } ) } else { @@ -180,7 +181,9 @@ flux_open_checkpoint <- function(transformer_dir) { } h <- safetensors::safetensors$new(single_path, framework = "torch") keys <- setdiff(h$keys(), "__metadata__") - handle <- list(get_tensor = function(key) h$get_tensor(key)) + handle <- list(get_tensor = function(key) { + .st_read_or_breadcrumb(function() h$get_tensor(key), single_path) + }) } list(handle = handle, keys = keys) } diff --git a/R/download_flux.R b/R/download_flux.R index 9c22c27..ed5a276 100644 --- a/R/download_flux.R +++ b/R/download_flux.R @@ -78,6 +78,9 @@ download_flux1 <- function(quantize = TRUE, precision = c("nf4", "fp8"), output_dir = NULL, text_encoders = TRUE, verbose = TRUE) { precision <- match.arg(precision) + # Explicit fp8 without float8 support: warn + build nf4 instead of + # failing in flux_quantize. + precision <- .st_graceful_precision(precision, mode = "write") if (is.null(output_dir)) { output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), paste0("flux1-schnell-", precision)) diff --git a/R/download_flux2.R b/R/download_flux2.R index a3a0fe2..2a1c155 100644 --- a/R/download_flux2.R +++ b/R/download_flux2.R @@ -51,6 +51,8 @@ download_flux2_klein <- function(quantize = TRUE, precision <- match.arg(precision) precision <- .flux_resolve_precision(precision, file.path(tools::R_user_dir("diffuseR", "data"), "flux2-klein-4b-")) + # Explicit fp8 without float8 support: warn + build nf4 rather than fail. + precision <- .st_graceful_precision(precision, mode = "write") if (is.null(output_dir)) { output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), paste0("flux2-klein-4b-", precision)) diff --git a/R/download_zimage.R b/R/download_zimage.R index fedfa22..78b6739 100644 --- a/R/download_zimage.R +++ b/R/download_zimage.R @@ -57,6 +57,8 @@ download_zimage_turbo <- function(quantize = TRUE, precision <- match.arg(precision) precision <- .flux_resolve_precision(precision, file.path(tools::R_user_dir("diffuseR", "data"), "zimage-turbo-")) + # Explicit fp8 without float8 support: warn + build nf4 rather than fail. + precision <- .st_graceful_precision(precision, mode = "write") if (is.null(output_dir)) { output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), paste0("zimage-turbo-", precision)) diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index a6e737d..ff9dad6 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -224,13 +224,17 @@ ltx23_open_fp8_checkpoint <- function(dir) { } manifest <- jsonlite::fromJSON(manifest_path, simplifyVector = TRUE) - handles <- lapply(file.path(dir, manifest$shards), function(p) { + shard_paths <- file.path(dir, manifest$shards) + handles <- lapply(shard_paths, function(p) { safetensors::safetensors$new(p, framework = "torch") }) key_to_handle <- list() - for (h in handles) { + key_to_path <- list() + for (i in seq_along(handles)) { + h <- handles[[i]] for (k in setdiff(h$keys(), "__metadata__")) { key_to_handle[[k]] <- h + key_to_path[[k]] <- shard_paths[[i]] } } @@ -238,7 +242,8 @@ ltx23_open_fp8_checkpoint <- function(dir) { get_tensor = function(key) { h <- key_to_handle[[key]] if (is.null(h)) stop("Key not found in fp8 shards: ", key) - h$get_tensor(key) + .st_read_or_breadcrumb(function() h$get_tensor(key), + key_to_path[[key]]) } ) diff --git a/R/st_caps.R b/R/st_caps.R index 3732fec..c4003a0 100644 --- a/R/st_caps.R +++ b/R/st_caps.R @@ -95,10 +95,11 @@ NULL } # The standard "install the fork, or press on with nf4" message. Shared -# by the recommender (read side) and the download graceful-fallback path -# (write side) so the wording stays identical everywhere. No em dashes -# (house style). -.st_fork_note <- function(precision) { +# by the recommender (read side, fit = TRUE: "best fit for your card") +# and the download graceful-fallback path (write side, fit = FALSE, since +# the user asked for it outright) so the wording stays identical +# everywhere. No em dashes (house style). +.st_fork_note <- function(precision, fit = TRUE) { precision <- as.character(precision) detail <- switch(precision, fp8 = "float8 support (mlverse/safetensors#13)", @@ -106,11 +107,77 @@ NULL bf16 = "bfloat16 write support (mlverse/safetensors#11)", bfloat16 = "bfloat16 write support (mlverse/safetensors#11)", paste0(precision, " support (mlverse/safetensors)")) + lead <- if (fit) { + sprintf("%s is the best fit for your card but needs", precision) + } else { + sprintf("%s needs", precision) + } sprintf(paste0( - "%s is the best fit for your card but needs ", - "cornball-ai/safetensors until CRAN safetensors ships %s. ", - "Install remotes::install_github(\"cornball-ai/safetensors\"), ", - "or press on with nf4: same weights, slightly lower ", - "precision, and it just works."), - precision, detail) + "%s cornball-ai/safetensors until CRAN safetensors ships ", + "%s. Install ", + "remotes::install_github(\"cornball-ai/safetensors\"), or ", + "press on with nf4: same weights, slightly lower precision, ", + "and it just works."), + lead, detail) +} + +# When a user explicitly asks for fp8/bf16 but the needed safetensors +# capability is missing, print the fork suggestion and fall back to nf4 +# instead of letting a downstream builder or loader fail. nf4, fp16, +# fp32 and anything unrecognized pass through untouched. `mode` selects +# the capability that matters: "write" when about to BUILD an artifact, +# "read" when about to LOAD one. +.st_graceful_precision <- function(precision, mode = c("write", "read"), + verbose = TRUE) { + mode <- match.arg(mode) + cap <- switch(precision, fp8 = "float8_e4m3fn", bf16 = "bfloat16", + bfloat16 = "bfloat16", NULL) + if (is.null(cap)) { + return(precision) + } + ok <- if (mode == "write") .st_can_write(cap) else .st_can_read(cap) + if (ok) { + return(precision) + } + if (verbose) { + message(.st_fork_note(precision, fit = FALSE), + "\nFalling back to nf4 for now.") + } + "nf4" +} + +# Actionable message for the >2 GB shard-read overflow. Split out from +# .st_read_or_breadcrumb so it can be unit-tested without a real 2 GB +# file. +.st_overflow_message <- function(file_path, size_bytes, underlying) { + sprintf(paste0( + "Could not read %s (%.1f GB). Stock CRAN safetensors ", + "overflows a 32-bit offset on files at or above 2^31 ", + "bytes (~2.15 GB). Rebuild the artifact with smaller ", + "shards (the quantizers now default to ", + "shard_bytes = 1.9e9), or install ", + "remotes::install_github(\"cornball-ai/safetensors\"). ", + "Underlying error: %s"), + basename(file_path), size_bytes / 1e9, underlying) +} + +# Run a safetensors read; if it fails AND the backing shard is at/above +# the 2^31-byte ceiling, translate the cryptic overflow into the +# fork-or-smaller-shards breadcrumb. A read that succeeds (fork, or a +# sub-2 GB shard) is untouched; a failure on a small shard rethrows +# verbatim. Reactive by design, so it never false-alarms on a machine +# that can read large files. +.st_read_or_breadcrumb <- function(read_fn, file_path = NULL) { + tryCatch(read_fn(), error = function(e) { + sz <- if (!is.null(file_path)) { + tryCatch(file.size(file_path), error = function(...) NA_real_) + } else { + NA_real_ + } + if (!is.na(sz) && sz >= 2^31) { + stop(.st_overflow_message(file_path, sz, conditionMessage(e)), + call. = FALSE) + } + stop(e) + }) } diff --git a/inst/tinytest/test_recommend.R b/inst/tinytest/test_recommend.R index 7611093..bf378f7 100644 --- a/inst/tinytest/test_recommend.R +++ b/inst/tinytest/test_recommend.R @@ -111,3 +111,48 @@ p2 <- flux_memory_profile(vram_gb = 16) expect_equal(p2$precision, "nf4") expect_true(isTRUE(p2$fork_suggested)) options(diffuseR.st_read_caps = NULL) + +# --- graceful fallback for an explicitly requested precision ---------------------- + +grc <- diffuseR:::.st_graceful_precision +# nf4/fp16 pass through untouched +expect_equal(grc("nf4", "write"), "nf4") +expect_equal(grc("fp16", "write"), "fp16") +# fp8 without float8 WRITE support -> nf4 + fork message +options(diffuseR.st_caps = list(float8_e4m3fn = FALSE)) +expect_message(rr <- grc("fp8", "write"), pattern = "safetensors#13") +expect_equal(rr, "nf4") +options(diffuseR.st_caps = list(float8_e4m3fn = TRUE)) +expect_equal(grc("fp8", "write"), "fp8") # present -> passes through +options(diffuseR.st_caps = NULL) +# bf16 gate goes through the READ probe in read mode +options(diffuseR.st_read_caps = list(bfloat16 = FALSE)) +expect_message(rr2 <- grc("bf16", "read"), pattern = "safetensors#11") +expect_equal(rr2, "nf4") +options(diffuseR.st_read_caps = NULL) + +# --- fork note fit parameter ------------------------------------------------------ + +fn <- diffuseR:::.st_fork_note +expect_true(grepl("best fit for your card", fn("fp8", fit = TRUE))) +expect_false(grepl("best fit for your card", fn("fp8", fit = FALSE))) +expect_false(grepl("—", fn("fp8"))) # no em dash, either variant + +# --- multi-GB read breadcrumb ----------------------------------------------------- + +msg <- diffuseR:::.st_overflow_message("shard-00001.safetensors", 3.4e9, "boom") +expect_true(grepl("3.4 GB", msg)) +expect_true(grepl("2\\^31", msg)) +expect_true(grepl("shard-00001", msg)) + +brc <- diffuseR:::.st_read_or_breadcrumb +# a read that succeeds is returned untouched +expect_equal(brc(function() 42L, NULL), 42L) +# a failure on a small (or absent) file rethrows verbatim, no breadcrumb +small <- tempfile() +writeLines("x", small) +e <- tryCatch(brc(function() stop("plain boom"), small), + error = function(err) conditionMessage(err)) +expect_true(grepl("plain boom", e)) +expect_false(grepl("2\\^31", e)) +unlink(small) From 60daa1af54b3c7ab5be3ca6cfade0162710d4028 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 14:53:48 -0500 Subject: [PATCH 05/18] NEWS + README: nf4 is CRAN-readable via sub-2 GB shards, not automatically Documents recommend(), the small-shard rationale, the graceful fallback, and the read breadcrumb. Corrects the impression that #28 alone made nf4 load on stock CRAN safetensors: the sub-2 GB shard default is what does it, since R safetensors overflows on files at/above 2^31 bytes. --- NEWS.md | 31 +++++++++++++++++++++++++++++++ README.md | 24 ++++++++++++++++++++++++ 2 files changed, 55 insertions(+) create mode 100644 NEWS.md diff --git a/NEWS.md b/NEWS.md new file mode 100644 index 0000000..4018227 --- /dev/null +++ b/NEWS.md @@ -0,0 +1,31 @@ +# diffuseR 0.1.0.6 (development) + +## Uniform native-safetensors + hosted quantization (in progress) + +* New `recommend(model, vram_gb, st_caps)`: one VRAM- and + capability-aware precision/device recommendation for every model + (sd21, sdxl, flux1, flux2, zimage, ltx). nf4 is the default tier; fp8 + or bf16 is recommended only when the card fits it *and* the installed + safetensors can **read** that dtype, otherwise it recommends nf4 and + surfaces the `cornball-ai/safetensors` suggestion in `$note` (never an + error). +* `flux_memory_profile()` now delegates to `recommend("flux1")`, + correcting the stale tiers that placed fp8 (GPU-resident now, not + streamed) in a narrow low-VRAM band it can no longer fit. +* Quantizer shards default to `shard_bytes = 1.9e9` (`flux_quantize`, + `ltx23_quantize_nf4`, `ltx23_quantize_fp8`). This is *what makes* nf4 + artifacts load on stock CRAN safetensors: R safetensors overflows a + 32-bit offset on any file at or above 2^31 bytes (~2.15 GB), so the + old 4e9 default produced shards only the fork could read. nf4 is + CRAN-readable **because of** the sub-2 GB shards, not automatically; + 4e9 remains available for local fork builds. +* Explicitly requesting fp8/bf16 without the needed safetensors support + now warns and falls back to nf4 instead of failing + (`download_flux1`, `download_flux2_klein`, `download_zimage_turbo`). +* Reading a legacy oversize (>2 GB) shard on stock safetensors raises an + actionable "rebuild with smaller shards or install the fork" message + instead of a raw 32-bit overflow error. + +All of the above is capability-**probed**, not version-pinned, so the +fork requirement self-heals when the safetensors fixes reach CRAN +(mlverse/safetensors#11, #13). diff --git a/README.md b/README.md index 77bef08..2c6687b 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,30 @@ txt2img_zimage(paste("A storefront with a large wooden sign that reads", txt2img("a lighthouse at dusk", model_name = "flux2") ``` +### Precision and safetensors + +The quantized transformers ship as **nf4** by default (packed uint8 + +float32 blocks). nf4 loads on **stock CRAN safetensors** because the +artifacts are written in sub-2 GB shards; R's safetensors overflows a +32-bit offset on any single file at or above 2^31 bytes (~2.15 GB), so +shard size, not the dtype, is what gates readability. + +Higher-quality tiers behave as follows: + +- **bf16** (24 GB+ cards) is also CRAN-readable. +- **fp8** (the 12-16 GB sweet spot) needs the + [`cornball-ai/safetensors`](https://github.com/cornball-ai/safetensors) + fork until float8 support lands on CRAN (mlverse/safetensors#13). + +`recommend(model)` picks the right tier for your VRAM and the +safetensors you have installed, and asking for a tier your safetensors +cannot read falls back to nf4 with a note rather than erroring: + +```r +recommend("flux2") # e.g. list(precision = "fp8", ...) on a 16 GB card +recommend("flux1")$note # the fork suggestion, when fp8/bf16 would fit but can't load +``` + ## Supported Models Currently supported models: From d979a8bc7bde885f0e3adde8382141d4d751bee9 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 15:35:00 -0500 Subject: [PATCH 06/18] Phase 1: safetensors weight loading for the native SD/SDXL UNet load_unet_safetensors / load_unet_sdxl_safetensors map the diffusers UNet2DConditionModel state-dict keys 1:1 into the native modules, the only rename being time_embedding.linear_{1,2} -> time_embedding_linear_ {1,2} (SDXL adds add_embedding.linear_{1,2}). Reuses .flux_open_sharded_dir so single-file and sharded checkpoints both work and the >2 GB read breadcrumb comes along; verifies every native parameter is filled and no key or shape is left unmatched. Validated against the real cached SDXL base UNet: all 1680 keys map with matching shapes. The VAE decoder and CLIP text encoder already had safetensors loaders; the UNet was the gap. Portable test drives the copy/completeness/error paths through a mock module. --- R/unet_safetensors.R | 122 ++++++++++++++++++++++++++ inst/tinytest/test_unet_safetensors.R | 116 ++++++++++++++++++++++++ 2 files changed, 238 insertions(+) create mode 100644 R/unet_safetensors.R create mode 100644 inst/tinytest/test_unet_safetensors.R diff --git a/R/unet_safetensors.R b/R/unet_safetensors.R new file mode 100644 index 0000000..43f18e2 --- /dev/null +++ b/R/unet_safetensors.R @@ -0,0 +1,122 @@ +#' Load HF safetensors weights into the native SD/SDXL UNet +#' +#' The native UNet modules mirror the diffusers +#' \code{UNet2DConditionModel} state-dict keys 1:1, with the sole +#' exception that the time- (and, for SDXL, add-) embedding MLPs are +#' flattened from dotted to underscored names +#' (\code{time_embedding.linear_1} -> \code{time_embedding_linear_1}). +#' These loaders read \code{unet/diffusion_pytorch_model.safetensors} +#' (single file or sharded via its \code{.index.json}) and copy each +#' weight into the matching native parameter, verifying that every native +#' parameter is filled and no key or shape is left unmatched. +#' +#' Reads route through the shared sharded opener, so an oversize (>2 GB) +#' single-file checkpoint on stock CRAN safetensors surfaces the +#' actionable "rebuild with smaller shards or install the fork" message +#' rather than a raw 32-bit overflow. +#' +#' @name unet_safetensors +NULL + +# time_embedding.linear_{1,2} -> time_embedding_linear_{1,2} +.unet_remap_sd21 <- function(key) { + key <- sub("^time_embedding\\.linear_1", "time_embedding_linear_1", key) + sub("^time_embedding\\.linear_2", "time_embedding_linear_2", key) +} + +# SD21 rules plus add_embedding.linear_{1,2} -> add_embedding_linear_{1,2} +.unet_remap_sdxl <- function(key) { + key <- .unet_remap_sd21(key) + key <- sub("^add_embedding\\.linear_1", "add_embedding_linear_1", key) + sub("^add_embedding\\.linear_2", "add_embedding_linear_2", key) +} + +# Shared loader: open the (single or sharded) diffusers UNet directory, +# map each HF key to a native parameter through `remap`, copy with a +# shape check, and fail loudly on any unmapped key, shape mismatch, or +# unfilled native parameter. +.load_unet_safetensors <- function(native_unet, path, remap, label, + verbose = TRUE) { + if (!requireNamespace("safetensors", quietly = TRUE)) { + stop("The safetensors package is required to read UNet weights.") + } + path <- path.expand(path) + dir <- if (dir.exists(path)) path else dirname(path) + opened <- .flux_open_sharded_dir(dir, "diffusion_pytorch_model") + keys <- opened$keys + + dests <- native_unet$named_parameters() + filled <- character(0) + unmapped <- character(0) + mismatch <- character(0) + + torch::with_no_grad({ + for (key in keys) { + native_name <- remap(key) + dest <- dests[[native_name]] + if (is.null(dest)) { + unmapped <- c(unmapped, key) + next + } + src <- opened$handle$get_tensor(key) + if (!all(dest$shape == src$shape)) { + mismatch <- c(mismatch, sprintf("%s (%s vs %s)", native_name, + paste(as.integer(src$shape), + collapse = "x"), + paste(as.integer(dest$shape), + collapse = "x"))) + next + } + dest$copy_(src) + filled <- c(filled, native_name) + } + }) + + if (length(unmapped)) { + stop(label, " load: ", length(unmapped), " unmapped keys, e.g. ", + paste(utils::head(unmapped, 3), collapse = ", ")) + } + if (length(mismatch)) { + stop(label, " load: ", length(mismatch), " shape mismatches, e.g. ", + paste(utils::head(mismatch, 3), collapse = ", ")) + } + unfilled <- setdiff(names(dests), filled) + if (length(unfilled)) { + stop(label, " load: ", length(unfilled), " unfilled params, e.g. ", + paste(utils::head(unfilled, 3), collapse = ", ")) + } + if (verbose) { + message("Loaded ", length(filled), " ", label, " parameters") + } + invisible(native_unet) +} + +#' Load HF safetensors weights into the native SD21 UNet +#' +#' @param native_unet A \code{\link{unet_native}} module. +#' @param path Path to the UNet directory (containing +#' \code{diffusion_pytorch_model.safetensors} or its shard index) or +#' directly to the single-file checkpoint. +#' @param verbose Print how many parameters were loaded. +#' +#' @return The native UNet with weights loaded (invisibly). +#' @export +load_unet_safetensors <- function(native_unet, path, verbose = TRUE) { + .load_unet_safetensors(native_unet, path, .unet_remap_sd21, "SD21 UNet", + verbose = verbose) +} + +#' Load HF safetensors weights into the native SDXL UNet +#' +#' @param native_unet A \code{\link{unet_sdxl_native}} module. +#' @param path Path to the UNet directory (containing +#' \code{diffusion_pytorch_model.safetensors} or its shard index) or +#' directly to the single-file checkpoint. +#' @param verbose Print how many parameters were loaded. +#' +#' @return The native UNet with weights loaded (invisibly). +#' @export +load_unet_sdxl_safetensors <- function(native_unet, path, verbose = TRUE) { + .load_unet_safetensors(native_unet, path, .unet_remap_sdxl, "SDXL UNet", + verbose = verbose) +} diff --git a/inst/tinytest/test_unet_safetensors.R b/inst/tinytest/test_unet_safetensors.R new file mode 100644 index 0000000..fcc1e8d --- /dev/null +++ b/inst/tinytest/test_unet_safetensors.R @@ -0,0 +1,116 @@ +# SD/SDXL UNet safetensors loader: HF-key remap rules, the copy path, +# completeness, and error reporting. The remap is validated against real +# cached SDXL weights out of band (all 1680 keys map with matching +# shapes); here we exercise the loader logic portably through a mock +# module so no multi-GB checkpoint is needed. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} +library(diffuseR) + +r21 <- diffuseR:::.unet_remap_sd21 +rxl <- diffuseR:::.unet_remap_sdxl + +# --- remap rules ------------------------------------------------------------------ + +expect_equal(r21("time_embedding.linear_1.weight"), + "time_embedding_linear_1.weight") +expect_equal(r21("time_embedding.linear_2.bias"), + "time_embedding_linear_2.bias") +# non-embedding keys pass through untouched (dotted block paths) +expect_equal(r21("down_blocks.0.resnets.0.norm1.weight"), + "down_blocks.0.resnets.0.norm1.weight") +expect_equal(r21("conv_out.weight"), "conv_out.weight") +# SDXL adds add_embedding and keeps the time_embedding rule +expect_equal(rxl("add_embedding.linear_1.weight"), + "add_embedding_linear_1.weight") +expect_equal(rxl("add_embedding.linear_2.bias"), "add_embedding_linear_2.bias") +expect_equal(rxl("time_embedding.linear_1.weight"), + "time_embedding_linear_1.weight") +# SD21 rule leaves add_embedding alone (SD21 has none) +expect_equal(r21("add_embedding.linear_1.weight"), + "add_embedding.linear_1.weight") + +# --- round-trip through a mock module --------------------------------------------- + +native_names <- c("conv_in.weight", "conv_in.bias", + "time_embedding_linear_1.weight", "time_embedding_linear_1.bias", + "time_embedding_linear_2.weight", + "down_blocks.0.resnets.0.norm1.weight", "add_embedding_linear_1.weight") + +# inverse remap (native -> HF) to author a synthetic checkpoint +to_hf <- function(n) { + n <- sub("^time_embedding_linear_1", "time_embedding.linear_1", n) + n <- sub("^time_embedding_linear_2", "time_embedding.linear_2", n) + sub("^add_embedding_linear_1", "add_embedding.linear_1", n) +} + +make_params <- function(gen) { + p <- lapply(native_names, function(i) gen()) + names(p) <- native_names + p +} +truth <- make_params(function() torch::torch_randn(2L, 3L)) + +dir <- tempfile() +dir.create(dir) +hf <- truth +names(hf) <- vapply(native_names, to_hf, character(1)) +safetensors::safe_save_file(hf, + file.path(dir, "diffusion_pytorch_model.safetensors")) + +# fresh zeroed params; the loader must copy truth into them +params <- make_params(function() torch::torch_zeros(2L, 3L)) +mock <- list(named_parameters = function() params) +diffuseR:::.load_unet_safetensors(mock, dir, rxl, "mock UNet", verbose = FALSE) +for (nm in native_names) { + expect_true(as.logical(torch::torch_allclose(params[[nm]], truth[[nm]]))) +} + +# passing the file directly (not the dir) also works +params_b <- make_params(function() torch::torch_zeros(2L, 3L)) +mock_b <- list(named_parameters = function() params_b) +diffuseR:::.load_unet_safetensors(mock_b, + file.path(dir, "diffusion_pytorch_model.safetensors"), rxl, "mock", + verbose = FALSE) +expect_true(as.logical(torch::torch_allclose(params_b[["conv_in.weight"]], + truth[["conv_in.weight"]]))) + +# --- error paths ------------------------------------------------------------------ + +# extra HF key with no native destination -> unmapped +dir2 <- tempfile() +dir.create(dir2) +hf_extra <- hf +hf_extra[["nonexistent.layer.weight"]] <- torch::torch_zeros(2L, 3L) +safetensors::safe_save_file(hf_extra, + file.path(dir2, "diffusion_pytorch_model.safetensors")) +expect_error( + diffuseR:::.load_unet_safetensors( + list(named_parameters = function() make_params(function() torch::torch_zeros(2L, 3L))), + dir2, rxl, "mock", verbose = FALSE), + pattern = "unmapped") + +# a native param with no checkpoint key -> unfilled +params_missing <- c(make_params(function() torch::torch_zeros(2L, 3L)), + list("extra.param.weight" = torch::torch_zeros(2L, 3L))) +expect_error( + diffuseR:::.load_unet_safetensors( + list(named_parameters = function() params_missing), dir, rxl, "mock", + verbose = FALSE), + pattern = "unfilled") + +# shape mismatch -> error (checked before unfilled) +params_wrong <- make_params(function() torch::torch_zeros(2L, 3L)) +params_wrong[["conv_in.weight"]] <- torch::torch_zeros(5L, 5L) +expect_error( + diffuseR:::.load_unet_safetensors( + list(named_parameters = function() params_wrong), dir, rxl, "mock", + verbose = FALSE), + pattern = "shape mismatch") + +unlink(c(dir, dir2), recursive = TRUE) From 1f0a62c489daebeeb087ddb7858d0a8514653b0a Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 15:35:00 -0500 Subject: [PATCH 07/18] document(): man + NAMESPACE for the SD/SDXL UNet safetensors loaders --- NAMESPACE | 2 ++ man/load_unet_safetensors.Rd | 22 ++++++++++++++++++++++ man/load_unet_sdxl_safetensors.Rd | 22 ++++++++++++++++++++++ man/unet_safetensors.Rd | 22 ++++++++++++++++++++++ 4 files changed, 68 insertions(+) create mode 100644 man/load_unet_safetensors.Rd create mode 100644 man/load_unet_sdxl_safetensors.Rd create mode 100644 man/unet_safetensors.Rd diff --git a/NAMESPACE b/NAMESPACE index ae1eda3..db039cd 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -80,6 +80,8 @@ export(load_text_encoder_safetensors) export(load_text_encoder_weights) export(load_text_encoder2_weights) export(load_to_gpu) +export(load_unet_safetensors) +export(load_unet_sdxl_safetensors) export(load_unet_sdxl_weights) export(load_unet_weights) export(ltx23_ada_layer_norm_single) diff --git a/man/load_unet_safetensors.Rd b/man/load_unet_safetensors.Rd new file mode 100644 index 0000000..bd34ed7 --- /dev/null +++ b/man/load_unet_safetensors.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_unet_safetensors} +\alias{load_unet_safetensors} +\title{Load HF safetensors weights into the native SD21 UNet} +\usage{ +load_unet_safetensors(native_unet, path, verbose = TRUE) +} +\arguments{ +\item{native_unet}{A \code{\link{unet_native}} module.} + +\item{path}{Path to the UNet directory (containing +\code{diffusion_pytorch_model.safetensors} or its shard index) or +directly to the single-file checkpoint.} + +\item{verbose}{Print how many parameters were loaded.} +} +\value{ +The native UNet with weights loaded (invisibly). +} +\description{ +Load HF safetensors weights into the native SD21 UNet +} diff --git a/man/load_unet_sdxl_safetensors.Rd b/man/load_unet_sdxl_safetensors.Rd new file mode 100644 index 0000000..7432490 --- /dev/null +++ b/man/load_unet_sdxl_safetensors.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{load_unet_sdxl_safetensors} +\alias{load_unet_sdxl_safetensors} +\title{Load HF safetensors weights into the native SDXL UNet} +\usage{ +load_unet_sdxl_safetensors(native_unet, path, verbose = TRUE) +} +\arguments{ +\item{native_unet}{A \code{\link{unet_sdxl_native}} module.} + +\item{path}{Path to the UNet directory (containing +\code{diffusion_pytorch_model.safetensors} or its shard index) or +directly to the single-file checkpoint.} + +\item{verbose}{Print how many parameters were loaded.} +} +\value{ +The native UNet with weights loaded (invisibly). +} +\description{ +Load HF safetensors weights into the native SDXL UNet +} diff --git a/man/unet_safetensors.Rd b/man/unet_safetensors.Rd new file mode 100644 index 0000000..649ec21 --- /dev/null +++ b/man/unet_safetensors.Rd @@ -0,0 +1,22 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{unet_safetensors} +\alias{unet_safetensors} +\title{Load HF safetensors weights into the native SD/SDXL UNet} +\description{ +The native UNet modules mirror the diffusers +\code{UNet2DConditionModel} state-dict keys 1:1, with the sole +exception that the time- (and, for SDXL, add-) embedding MLPs are +flattened from dotted to underscored names +(\code{time_embedding.linear_1} -> \code{time_embedding_linear_1}). +These loaders read \code{unet/diffusion_pytorch_model.safetensors} +(single file or sharded via its \code{.index.json}) and copy each +weight into the matching native parameter, verifying that every native +parameter is filled and no key or shape is left unmatched. +} +\details{ +Reads route through the shared sharded opener, so an oversize (>2 GB) +single-file checkpoint on stock CRAN safetensors surfaces the +actionable "rebuild with smaller shards or install the fork" message +rather than a raw 32-bit overflow. + +} From e07ab32f01fd6f375bf4290460ebdb2f5e8ee4da Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 15:35:00 -0500 Subject: [PATCH 08/18] Bump version to 0.1.0.6 --- DESCRIPTION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index cd707a3..b5cb6a9 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.1.0.5 +Version: 0.1.0.6 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), From c14071f5b992144c96e93b5a57719cd33bdac4d6 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 16:30:41 -0500 Subject: [PATCH 09/18] Native SD/SDXL UNet constructors from safetensors unet_native_from_safetensors / unet_sdxl_native_from_safetensors build the native UNet and load its weights from a diffusers directory, the safetensors counterparts to the *_from_torchscript constructors (no TorchScript, so Blackwell-safe). Thin wrappers over the validated default construction plus the safetensors loader; ... forwards variant overrides and the loader's shape check fails loudly on a mismatched architecture. Validated end-to-end against the cached SDXL base UNet: 1680 params built + loaded in ~20 s. Opt-in gated test documents it. --- NEWS.md | 7 ++++ R/unet_safetensors.R | 48 +++++++++++++++++++++++++++ inst/tinytest/test_unet_safetensors.R | 17 ++++++++++ 3 files changed, 72 insertions(+) diff --git a/NEWS.md b/NEWS.md index 4018227..7f61ce1 100644 --- a/NEWS.md +++ b/NEWS.md @@ -25,6 +25,13 @@ * Reading a legacy oversize (>2 GB) shard on stock safetensors raises an actionable "rebuild with smaller shards or install the fork" message instead of a raw 32-bit overflow error. +* Native SD21/SDXL UNet weights now load from diffusers safetensors + (`load_unet_safetensors`, `load_unet_sdxl_safetensors`, and the + `unet_native_from_safetensors` / `unet_sdxl_native_from_safetensors` + constructors), with no TorchScript step (Blackwell-safe). The VAE + decoder and CLIP text encoder already had safetensors loaders; the + UNet was the gap. Validated against the cached SDXL base UNet (all + 1680 keys map with matching shapes). All of the above is capability-**probed**, not version-pinned, so the fork requirement self-heals when the safetensors fixes reach CRAN diff --git a/R/unet_safetensors.R b/R/unet_safetensors.R index 43f18e2..9667eb5 100644 --- a/R/unet_safetensors.R +++ b/R/unet_safetensors.R @@ -120,3 +120,51 @@ load_unet_sdxl_safetensors <- function(native_unet, path, verbose = TRUE) { .load_unet_safetensors(native_unet, path, .unet_remap_sdxl, "SDXL UNet", verbose = verbose) } + +#' Build a native SD21 UNet from a diffusers safetensors directory +#' +#' The safetensors counterpart to +#' \code{\link{unet_native_from_torchscript}}: constructs +#' \code{\link{unet_native}} and loads its weights from +#' \code{unet/diffusion_pytorch_model.safetensors} (no TorchScript, so it +#' works on Blackwell). The default construction matches the canonical +#' Stable Diffusion 2.1 UNet; pass constructor overrides through +#' \code{...} for a variant checkpoint (the loader fails loudly on any +#' shape mismatch, so a wrong architecture surfaces immediately rather +#' than loading silently wrong weights). +#' +#' @param path Path to the UNet directory or its single-file checkpoint. +#' @param verbose Print how many parameters were loaded. +#' @param ... Overrides for \code{\link{unet_native}} constructor args. +#' +#' @return The native SD21 UNet in eval mode. +#' @export +unet_native_from_safetensors <- function(path, verbose = TRUE, ...) { + model <- unet_native(...) + load_unet_safetensors(model, path, verbose = verbose) + model$eval() + model +} + +#' Build a native SDXL UNet from a diffusers safetensors directory +#' +#' The safetensors counterpart to +#' \code{\link{unet_sdxl_native_from_torchscript}}: constructs +#' \code{\link{unet_sdxl_native}} and loads its weights from +#' \code{unet/diffusion_pytorch_model.safetensors}. Validated against the +#' cached \code{stabilityai/stable-diffusion-xl-base-1.0} UNet (all 1680 +#' parameters map with matching shapes). Pass constructor overrides +#' through \code{...} for a variant checkpoint. +#' +#' @param path Path to the UNet directory or its single-file checkpoint. +#' @param verbose Print how many parameters were loaded. +#' @param ... Overrides for \code{\link{unet_sdxl_native}} constructor args. +#' +#' @return The native SDXL UNet in eval mode. +#' @export +unet_sdxl_native_from_safetensors <- function(path, verbose = TRUE, ...) { + model <- unet_sdxl_native(...) + load_unet_sdxl_safetensors(model, path, verbose = verbose) + model$eval() + model +} diff --git a/inst/tinytest/test_unet_safetensors.R b/inst/tinytest/test_unet_safetensors.R index fcc1e8d..79106c2 100644 --- a/inst/tinytest/test_unet_safetensors.R +++ b/inst/tinytest/test_unet_safetensors.R @@ -114,3 +114,20 @@ expect_error( pattern = "shape mismatch") unlink(c(dir, dir2), recursive = TRUE) + +# --- end-to-end against real SDXL weights (opt-in; heavy ~9.6 GB load) ------------- +# Runs only with DIFFUSER_TEST_SDXL_LOAD=1 and the cached diffusers UNet +# present, so a normal suite run never pulls 9.6 GB. Validated manually: +# builds + loads all 1680 params in ~20 s. +sdxl_dir <- Sys.glob(file.path("~/.cache/huggingface/hub", + "models--stabilityai--stable-diffusion-xl-base-1.0/snapshots/*/unet")) +if (at_home() && nzchar(Sys.getenv("DIFFUSER_TEST_SDXL_LOAD")) && + length(sdxl_dir) && dir.exists(sdxl_dir[1])) { + m <- unet_sdxl_native_from_safetensors(sdxl_dir[1], verbose = FALSE) + expect_equal(length(m$named_parameters()), 1680L) + w <- m$named_parameters()[["conv_in.weight"]] + expect_true(as.numeric(w$abs()$sum()$item()) > 0) + expect_false(m$training) + rm(m) + gc() +} From 0e50c081d5a608e9b196ad1bf3588c6d28657b0a Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 16:30:41 -0500 Subject: [PATCH 10/18] document(): man + NAMESPACE for the from-safetensors UNet constructors --- NAMESPACE | 2 ++ man/unet_native_from_safetensors.Rd | 28 ++++++++++++++++++++++++ man/unet_sdxl_native_from_safetensors.Rd | 26 ++++++++++++++++++++++ 3 files changed, 56 insertions(+) create mode 100644 man/unet_native_from_safetensors.Rd create mode 100644 man/unet_sdxl_native_from_safetensors.Rd diff --git a/NAMESPACE b/NAMESPACE index db039cd..0c1c677 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -190,8 +190,10 @@ export(txt2img_sdxl) export(txt2img_zimage) export(txt2vid_ltx2) export(unet_native) +export(unet_native_from_safetensors) export(unet_native_from_torchscript) export(unet_sdxl_native) +export(unet_sdxl_native_from_safetensors) export(unet_sdxl_native_from_torchscript) export(unigram_tokenizer) export(vae_decoder_native) diff --git a/man/unet_native_from_safetensors.Rd b/man/unet_native_from_safetensors.Rd new file mode 100644 index 0000000..60dc841 --- /dev/null +++ b/man/unet_native_from_safetensors.Rd @@ -0,0 +1,28 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{unet_native_from_safetensors} +\alias{unet_native_from_safetensors} +\title{Build a native SD21 UNet from a diffusers safetensors directory} +\usage{ +unet_native_from_safetensors(path, verbose = TRUE, ...) +} +\arguments{ +\item{path}{Path to the UNet directory or its single-file checkpoint.} + +\item{verbose}{Print how many parameters were loaded.} + +\item{...}{Overrides for \code{\link{unet_native}} constructor args.} +} +\value{ +The native SD21 UNet in eval mode. +} +\description{ +The safetensors counterpart to +\code{\link{unet_native_from_torchscript}}: constructs +\code{\link{unet_native}} and loads its weights from +\code{unet/diffusion_pytorch_model.safetensors} (no TorchScript, so it +works on Blackwell). The default construction matches the canonical +Stable Diffusion 2.1 UNet; pass constructor overrides through +\code{...} for a variant checkpoint (the loader fails loudly on any +shape mismatch, so a wrong architecture surfaces immediately rather +than loading silently wrong weights). +} diff --git a/man/unet_sdxl_native_from_safetensors.Rd b/man/unet_sdxl_native_from_safetensors.Rd new file mode 100644 index 0000000..f4a1b2d --- /dev/null +++ b/man/unet_sdxl_native_from_safetensors.Rd @@ -0,0 +1,26 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{unet_sdxl_native_from_safetensors} +\alias{unet_sdxl_native_from_safetensors} +\title{Build a native SDXL UNet from a diffusers safetensors directory} +\usage{ +unet_sdxl_native_from_safetensors(path, verbose = TRUE, ...) +} +\arguments{ +\item{path}{Path to the UNet directory or its single-file checkpoint.} + +\item{verbose}{Print how many parameters were loaded.} + +\item{...}{Overrides for \code{\link{unet_sdxl_native}} constructor args.} +} +\value{ +The native SDXL UNet in eval mode. +} +\description{ +The safetensors counterpart to +\code{\link{unet_sdxl_native_from_torchscript}}: constructs +\code{\link{unet_sdxl_native}} and loads its weights from +\code{unet/diffusion_pytorch_model.safetensors}. Validated against the +cached \code{stabilityai/stable-diffusion-xl-base-1.0} UNet (all 1680 +parameters map with matching shapes). Pass constructor overrides +through \code{...} for a variant checkpoint. +} From 8a9772a0121bf9e89be0cb0174d78ba1da6471b4 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 17:12:35 -0500 Subject: [PATCH 11/18] Native VAE-decoder + CLIP text-encoder constructors from safetensors vae_decoder_native_from_safetensors and text_encoder_native_from_safetensors build and load their native modules from a diffusers directory (no TorchScript, Blackwell-safe). The text-encoder builder reads CLIPTextConfig from config.json via .detect_text_encoder_config (the config counterpart to the TorchScript arch detector). Validated against cached FLUX analogs: the FLUX VAE exercises the shared 16-channel decoder, and FLUX's text_encoder IS the SDXL CLIP ViT-L (config-detected 768/12/12, forward -> 1x77x768). --- R/text_encoder.R | 56 +++++++++++++++++++++++ R/vae_decoder.R | 25 +++++++++++ inst/tinytest/test_sd_safetensors.R | 69 +++++++++++++++++++++++++++++ 3 files changed, 150 insertions(+) create mode 100644 inst/tinytest/test_sd_safetensors.R diff --git a/R/text_encoder.R b/R/text_encoder.R index 60687fc..2093a9e 100644 --- a/R/text_encoder.R +++ b/R/text_encoder.R @@ -386,6 +386,62 @@ load_text_encoder_safetensors <- function(native_encoder, path, invisible(native_encoder) } +# Detect a CLIP text encoder's architecture from a diffusers +# config.json (the config counterpart to +# detect_text_encoder_architecture, which reads a TorchScript file). +# Accepts a directory or a config.json path. +.detect_text_encoder_config <- function(path) { + cfg_path <- if (dir.exists(path)) file.path(path, "config.json") else path + if (!file.exists(cfg_path)) { + stop("No config.json for the text encoder at ", path) + } + cfg <- jsonlite::fromJSON(cfg_path, simplifyVector = TRUE) + list( + vocab_size = as.integer(cfg$vocab_size), + context_length = as.integer(cfg$max_position_embeddings), + embed_dim = as.integer(cfg$hidden_size), + num_layers = as.integer(cfg$num_hidden_layers), + num_heads = as.integer(cfg$num_attention_heads), + mlp_dim = as.integer(cfg$intermediate_size) + ) +} + +#' Build a native CLIP text encoder from a diffusers safetensors directory +#' +#' Reads the CLIPTextConfig from \code{/config.json}, constructs +#' \code{\link{text_encoder_native}} to match, and loads +#' \code{model.safetensors} - the safetensors counterpart to the +#' TorchScript text-encoder path (no TorchScript, Blackwell-safe). +#' Handles SD21's OpenCLIP ViT-H and SDXL's CLIP ViT-L (which is the same +#' checkpoint as FLUX's \code{text_encoder}). \code{apply_final_ln} +#' governs only the forward output; the \code{final_layer_norm} weights +#' load either way. Use \code{TRUE} for SD21 and pooled CLIP outputs, +#' \code{FALSE} for the SDXL penultimate-layer prompt embeds. +#' +#' @param path diffusers text_encoder directory (config.json + +#' model.safetensors) or the config.json path. +#' @param apply_final_ln Apply the final layer norm in forward (default TRUE). +#' @param verbose Print how many parameters were loaded. +#' @param ... Overrides for \code{\link{text_encoder_native}} args (e.g. +#' \code{gelu_type}). +#' +#' @return The native text encoder in eval mode. +#' @export +text_encoder_native_from_safetensors <- function(path, apply_final_ln = TRUE, + verbose = TRUE, ...) { + arch <- .detect_text_encoder_config(path) + args <- list(vocab_size = arch$vocab_size, + context_length = arch$context_length, + embed_dim = arch$embed_dim, num_layers = arch$num_layers, + num_heads = arch$num_heads, mlp_dim = arch$mlp_dim, + apply_final_ln = apply_final_ln) + args <- utils::modifyList(args, list(...)) + model <- do.call(text_encoder_native, args) + load_text_encoder_safetensors(model, path, verbose = verbose) + model$eval() + model +} + #' Load weights from TorchScript text encoder into native encoder #' #' @param native_encoder Native text encoder module diff --git a/R/vae_decoder.R b/R/vae_decoder.R index 73d76b4..ad13c48 100644 --- a/R/vae_decoder.R +++ b/R/vae_decoder.R @@ -231,6 +231,31 @@ load_decoder_safetensors <- function(native_decoder, path, verbose = TRUE) { invisible(native_decoder) } +#' Build a native VAE decoder from a diffusers safetensors directory +#' +#' The safetensors counterpart to the TorchScript decoder path: +#' constructs \code{\link{vae_decoder_native}} and loads the decoder half +#' of a diffusers AutoencoderKL checkpoint (no TorchScript, so it works +#' on Blackwell). \code{latent_channels} defaults to 4 (SD/SDXL); pass 16 +#' for the FLUX/SD3 VAE. The SD/SDXL and FLUX VAEs share the decoder +#' shape and differ only in that channel count. +#' +#' @param path Path to the VAE directory (containing +#' \code{diffusion_pytorch_model.safetensors}) or the file itself. +#' @param latent_channels Latent channel count (4 for SD/SDXL, 16 for FLUX). +#' @param verbose Print how many parameters were loaded. +#' @param ... Overrides for \code{\link{vae_decoder_native}} constructor args. +#' +#' @return The native VAE decoder in eval mode. +#' @export +vae_decoder_native_from_safetensors <- function(path, latent_channels = 4L, + verbose = TRUE, ...) { + model <- vae_decoder_native(latent_channels = latent_channels, ...) + load_decoder_safetensors(model, path, verbose = verbose) + model$eval() + model +} + #' Load weights from TorchScript decoder into native decoder #' #' @param native_decoder Native VAE decoder module diff --git a/inst/tinytest/test_sd_safetensors.R b/inst/tinytest/test_sd_safetensors.R new file mode 100644 index 0000000..7e8fa26 --- /dev/null +++ b/inst/tinytest/test_sd_safetensors.R @@ -0,0 +1,69 @@ +# SD-from-safetensors component constructors: config-based CLIP arch +# detection (portable), and the VAE-decoder + CLIP text-encoder +# from-safetensors builders (validated against cached FLUX analogs - +# FLUX's VAE exercises the shared decoder, and FLUX's text_encoder IS +# the SDXL CLIP ViT-L). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +if (!requireNamespace("safetensors", quietly = TRUE)) { + exit_file("safetensors not installed") +} +library(diffuseR) + +# --- config-based CLIP arch detection (portable) ---------------------------------- + +tmp <- tempfile() +dir.create(tmp) +writeLines(jsonlite::toJSON(list(vocab_size = 49408L, + max_position_embeddings = 77L, hidden_size = 1280L, + num_hidden_layers = 32L, num_attention_heads = 20L, + intermediate_size = 5120L), auto_unbox = TRUE), + file.path(tmp, "config.json")) +arch <- diffuseR:::.detect_text_encoder_config(tmp) +expect_equal(arch$vocab_size, 49408L) +expect_equal(arch$context_length, 77L) +expect_equal(arch$embed_dim, 1280L) # bigG dims +expect_equal(arch$num_layers, 32L) +expect_equal(arch$num_heads, 20L) +expect_equal(arch$mlp_dim, 5120L) +# accepts a config.json path directly too +expect_equal(diffuseR:::.detect_text_encoder_config( + file.path(tmp, "config.json"))$embed_dim, 1280L) +expect_error(diffuseR:::.detect_text_encoder_config(tempfile()), + pattern = "config.json") +unlink(tmp, recursive = TRUE) + +# --- against cached FLUX analogs (skipped where the cache is absent) --------------- + +flux_vae <- Sys.glob(file.path("~/.cache/huggingface/hub", + "models--black-forest-labs--FLUX.1-schnell/snapshots/*/vae")) +flux_te <- Sys.glob(file.path("~/.cache/huggingface/hub", + "models--black-forest-labs--FLUX.1-schnell/snapshots/*/text_encoder")) + +if (at_home() && length(flux_vae) && dir.exists(flux_vae[1])) { + d <- vae_decoder_native_from_safetensors(flux_vae[1], latent_channels = 16L, + verbose = FALSE) + ci <- d$named_parameters()[["conv_in.weight"]] + expect_equal(as.integer(ci$shape[2]), 16L) # 16-channel latent + expect_true(as.numeric(ci$abs()$sum()$item()) > 0) # actually loaded + expect_false(d$training) + rm(d) + gc() +} + +if (at_home() && length(flux_te) && dir.exists(flux_te[1])) { + a <- diffuseR:::.detect_text_encoder_config(flux_te[1]) + expect_equal(a$embed_dim, 768L) # CLIP ViT-L + expect_equal(a$num_layers, 12L) + enc <- text_encoder_native_from_safetensors(flux_te[1], verbose = FALSE) + expect_false(enc$training) + tok <- torch::torch_tensor( + matrix(c(49406L, 320L, 49407L, rep(49407L, 74)), nrow = 1), + dtype = torch::torch_long()) + out <- torch::with_no_grad(enc(tok)) + expect_equal(as.integer(out$shape), c(1L, 77L, 768L)) + rm(enc) + gc() +} From acd9be80b20a056abc18cf456c1cd38948d5b925 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Thu, 9 Jul 2026 17:12:35 -0500 Subject: [PATCH 12/18] document(): man + NAMESPACE for the SD component from-safetensors constructors --- NAMESPACE | 2 ++ man/text_encoder_native_from_safetensors.Rd | 33 +++++++++++++++++++++ man/vae_decoder_native_from_safetensors.Rd | 29 ++++++++++++++++++ 3 files changed, 64 insertions(+) create mode 100644 man/text_encoder_native_from_safetensors.Rd create mode 100644 man/vae_decoder_native_from_safetensors.Rd diff --git a/NAMESPACE b/NAMESPACE index 0c1c677..9137c0e 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -180,6 +180,7 @@ export(scheduler_add_noise) export(sdxl_memory_profile) export(t5_encoder) export(text_encoder_native) +export(text_encoder_native_from_safetensors) export(text_encoder2_native) export(tokenize_gemma3) export(txt2img) @@ -197,6 +198,7 @@ export(unet_sdxl_native_from_safetensors) export(unet_sdxl_native_from_torchscript) export(unigram_tokenizer) export(vae_decoder_native) +export(vae_decoder_native_from_safetensors) export(vocab_size) export(vram_report) export(write_wav) diff --git a/man/text_encoder_native_from_safetensors.Rd b/man/text_encoder_native_from_safetensors.Rd new file mode 100644 index 0000000..35450d0 --- /dev/null +++ b/man/text_encoder_native_from_safetensors.Rd @@ -0,0 +1,33 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{text_encoder_native_from_safetensors} +\alias{text_encoder_native_from_safetensors} +\title{Build a native CLIP text encoder from a diffusers safetensors directory} +\usage{ +text_encoder_native_from_safetensors(path, apply_final_ln = TRUE, + verbose = TRUE, ...) +} +\arguments{ +\item{path}{diffusers text_encoder directory (config.json + +model.safetensors) or the config.json path.} + +\item{apply_final_ln}{Apply the final layer norm in forward (default TRUE).} + +\item{verbose}{Print how many parameters were loaded.} + +\item{...}{Overrides for \code{\link{text_encoder_native}} args (e.g. +\code{gelu_type}).} +} +\value{ +The native text encoder in eval mode. +} +\description{ +Reads the CLIPTextConfig from \code{/config.json}, constructs +\code{\link{text_encoder_native}} to match, and loads +\code{model.safetensors} - the safetensors counterpart to the +TorchScript text-encoder path (no TorchScript, Blackwell-safe). +Handles SD21's OpenCLIP ViT-H and SDXL's CLIP ViT-L (which is the same +checkpoint as FLUX's \code{text_encoder}). \code{apply_final_ln} +governs only the forward output; the \code{final_layer_norm} weights +load either way. Use \code{TRUE} for SD21 and pooled CLIP outputs, +\code{FALSE} for the SDXL penultimate-layer prompt embeds. +} diff --git a/man/vae_decoder_native_from_safetensors.Rd b/man/vae_decoder_native_from_safetensors.Rd new file mode 100644 index 0000000..c1de476 --- /dev/null +++ b/man/vae_decoder_native_from_safetensors.Rd @@ -0,0 +1,29 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{vae_decoder_native_from_safetensors} +\alias{vae_decoder_native_from_safetensors} +\title{Build a native VAE decoder from a diffusers safetensors directory} +\usage{ +vae_decoder_native_from_safetensors(path, latent_channels = 4L, verbose = TRUE, + ...) +} +\arguments{ +\item{path}{Path to the VAE directory (containing +\code{diffusion_pytorch_model.safetensors}) or the file itself.} + +\item{latent_channels}{Latent channel count (4 for SD/SDXL, 16 for FLUX).} + +\item{verbose}{Print how many parameters were loaded.} + +\item{...}{Overrides for \code{\link{vae_decoder_native}} constructor args.} +} +\value{ +The native VAE decoder in eval mode. +} +\description{ +The safetensors counterpart to the TorchScript decoder path: +constructs \code{\link{vae_decoder_native}} and loads the decoder half +of a diffusers AutoencoderKL checkpoint (no TorchScript, so it works +on Blackwell). \code{latent_channels} defaults to 4 (SD/SDXL); pass 16 +for the FLUX/SD3 VAE. The SD/SDXL and FLUX VAEs share the decoder +shape and differ only in that channel count. +} From 40016c14b1b051f462d9f08d2643ef1beb3137ff Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 07:47:36 -0500 Subject: [PATCH 13/18] Native SD 2.1 pipeline from diffusers safetensors + VAE post_quant_conv download_sd21() fetches the SD 2.1 diffusers weights (gated-repo breadcrumb if the license is unaccepted); sd_pipeline_from_safetensors() assembles the native UNet + VAE decode + CLIP text encoder from a diffusers directory, and txt2img_sd21(diffusers_dir=) runs it with no TorchScript (Blackwell-safe). Fixes the SD VAE decode: the native decoder (derived from the FLUX VAE, which has no quant convs) omitted post_quant_conv, so decode(z) skipped the 1x1 conv diffusers applies before the decoder -- cos ~0.5 vs the reference, i.e. garbage. sd_vae_decode now applies it from the artifact's own post_quant_conv weights (cos 1.0). SD 2.1 defaults to float32 here (fp16 attention overflows to NaN). Known issue (pre-existing, tracked): the native SD 2.1 UNet has a spatial-coherence "tiling" bug that the .pt-native path shares (identical weights); the existing parity test misses it by feeding constant inputs. --- NEWS.md | 11 +++ R/sd_pipeline_safetensors.R | 191 ++++++++++++++++++++++++++++++++++++ R/txt2img_sd21.R | 53 +++++++--- 3 files changed, 242 insertions(+), 13 deletions(-) create mode 100644 R/sd_pipeline_safetensors.R diff --git a/NEWS.md b/NEWS.md index 7f61ce1..65a522d 100644 --- a/NEWS.md +++ b/NEWS.md @@ -32,6 +32,17 @@ decoder and CLIP text encoder already had safetensors loaders; the UNet was the gap. Validated against the cached SDXL base UNet (all 1680 keys map with matching shapes). +* `vae_decoder_native_from_safetensors` and + `text_encoder_native_from_safetensors` (config-driven CLIP arch + detection) complete the native SD component set. +* `download_sd21()` + `sd_pipeline_from_safetensors()` run Stable + Diffusion 2.1 fully natively from diffusers safetensors; + `txt2img_sd21(diffusers_dir=)` uses it. The SD VAE decode now applies + the `post_quant_conv` the FLUX-derived native decoder omitted (the + decode was badly wrong without it). SD 2.1 defaults to float32 on this + path (fp16 attention overflows to NaN). A pre-existing spatial-coherence + ("tiling") bug in the native SD 2.1 UNet is tracked separately; it + affects the `.pt`-native path identically (same weights). All of the above is capability-**probed**, not version-pinned, so the fork requirement self-heals when the safetensors fixes reach CRAN diff --git a/R/sd_pipeline_safetensors.R b/R/sd_pipeline_safetensors.R new file mode 100644 index 0000000..374ace7 --- /dev/null +++ b/R/sd_pipeline_safetensors.R @@ -0,0 +1,191 @@ +#' Native Stable Diffusion pipelines from diffusers safetensors +#' +#' Assemble and run the native SD pipeline directly from a HuggingFace +#' diffusers directory (\code{unet/}, \code{vae/}, \code{text_encoder/}), +#' with no TorchScript \code{.pt} step - so it works on Blackwell and +#' loads the same weights everyone else uses. SD21 is wired end to end +#' here; SDXL still needs its second text encoder and added-conditioning +#' embeddings (tracked in tasks/todo.md). +#' +#' @name sd_pipeline_safetensors +NULL + +.sd21_repo <- "stabilityai/stable-diffusion-2-1" + +# The SD/SDXL AutoencoderKL applies a 1x1 post_quant_conv to the latent +# before the decoder (decode(z) = decoder(post_quant_conv(z))). The +# native vae_decoder_native is the decoder submodule only (the FLUX VAE +# has no post_quant_conv), so the SD decode path must apply it - without +# it the decode is badly wrong (cos ~0.5 vs the reference). This wrapper +# restores it; its weights come from the artifact's own +# post_quant_conv.{weight,bias}. +sd_vae_decode <- torch::nn_module( + "SDVAEDecode", + initialize = function(decoder, post_quant_conv) { + self$post_quant_conv <- post_quant_conv + self$decoder <- decoder +}, + forward = function(z) { + self$decoder(self$post_quant_conv(z)) +} +) + +# Build the SD decode module (post_quant_conv + native decoder) from a +# diffusers VAE directory. +.sd_vae_decode_from_safetensors <- function(vae_dir, latent_channels = 4L) { + decoder <- vae_decoder_native_from_safetensors(vae_dir, + latent_channels = latent_channels, verbose = FALSE) + path <- file.path(path.expand(vae_dir), + "diffusion_pytorch_model.safetensors") + h <- safetensors::safetensors$new(path, framework = "torch") + pqc <- torch::nn_conv2d(latent_channels, latent_channels, kernel_size = 1L) + torch::with_no_grad({ + pqc$weight$copy_(h$get_tensor("post_quant_conv.weight")) + pqc$bias$copy_(h$get_tensor("post_quant_conv.bias")) + }) + m <- sd_vae_decode(decoder, pqc) + m$eval() + m +} + +.sd21_files <- c("unet/config.json", + "unet/diffusion_pytorch_model.safetensors", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", + "text_encoder/config.json", + "text_encoder/model.safetensors") + +#' Download the Stable Diffusion 2.1 diffusers weights +#' +#' Fetches the UNet, VAE, and CLIP text encoder (config + safetensors) +#' from \code{stabilityai/stable-diffusion-2-1} into the hfhub cache. The +#' native tokenizer and DDIM scheduler need no downloads. About 5 GB +#' (fp32 UNet); one-time. +#' +#' @param verbose Logical. +#' +#' @return Invisibly, the diffusers directory (the parent of +#' \code{unet/}, \code{vae/}, \code{text_encoder/}). +#' +#' @export +download_sd21 <- function(verbose = TRUE) { + if (!requireNamespace("hfhub", quietly = TRUE)) { + stop("The hfhub package is required to download model weights.") + } + have <- !is.null(tryCatch( + hfhub::hub_download(.sd21_repo, .sd21_files[[2]], + local_files_only = TRUE), + error = function(e) NULL + )) + if (!have) { + ok <- .ltx23_consent( + "Stable Diffusion 2.1 UNet + VAE + CLIP text encoder (~5 GB)" + ) + if (!ok) { + stop("Download cancelled.", call. = FALSE) + } + if (verbose) { + message("Downloading Stable Diffusion 2.1 (diffusers safetensors)...") + } + } + dl <- function(f) { + tryCatch(hfhub::hub_download(.sd21_repo, f), error = function(e) { + msg <- conditionMessage(e) + if (grepl("404|not found|[Uu]nauthorized|[Ff]orbidden|401|403", + msg)) { + stop( + "Could not fetch ", .sd21_repo, " (", sub("\n.*", "", msg), + "). This repo is gated on HuggingFace; a 404 with a valid ", + "token means the license has not been accepted yet. To fix:\n", + " 1. Visit https://huggingface.co/", .sd21_repo, + " and accept the license.\n", + " 2. Ensure HF_TOKEN is set (Sys.getenv(\"HF_TOKEN\")).\n", + "Or pass an accessible diffusers directory straight to ", + "sd_pipeline_from_safetensors() / txt2img_sd21(diffusers_dir=).", + call. = FALSE) + } + stop(e) + }) + } + paths <- vapply(.sd21_files, dl, character(1)) + # the diffusers root is the parent of unet/ (paths[[1]] is unet/config.json) + diffusers_dir <- dirname(dirname(paths[[1]])) + if (verbose) { + message("SD 2.1 ready: ", diffusers_dir) + } + invisible(diffusers_dir) +} + +#' Assemble a native SD pipeline from a diffusers safetensors directory +#' +#' Builds the native UNet, VAE decoder, and CLIP text encoder from a +#' diffusers directory using the \code{*_from_safetensors} constructors, +#' places each on its component device, and returns the \code{$unet / +#' $decoder / $text_encoder} list the \code{txt2img_*} denoise loop +#' expects. +#' +#' @param diffusers_dir Directory with \code{unet/}, \code{vae/}, +#' \code{text_encoder/} subdirectories. +#' @param model_name Currently "sd21" (SDXL pending its second encoder). +#' @param devices Named list of component devices (\code{unet}, +#' \code{decoder}, \code{text_encoder}); defaults to all-CPU. +#' @param unet_dtype A torch dtype for the UNet (default float16 on CUDA, +#' float32 on CPU). +#' @param verbose Logical. +#' +#' @return A list with \code{unet}, \code{decoder}, \code{text_encoder}. +#' +#' @export +sd_pipeline_from_safetensors <- function(diffusers_dir, model_name = "sd21", + devices = NULL, unet_dtype = NULL, + verbose = TRUE) { + if (!identical(model_name, "sd21")) { + stop("sd_pipeline_from_safetensors currently supports \"sd21\"; ", + "SDXL needs its second text encoder and added-conditioning ", + "embeddings (see tasks/todo.md).") + } + diffusers_dir <- path.expand(diffusers_dir) + if (is.null(devices)) { + devices <- list(unet = "cpu", decoder = "cpu", text_encoder = "cpu") + } + if (is.null(unet_dtype)) { + unet_dtype <- if (identical(devices$unet, "cuda")) { + torch::torch_float16() + } else { + torch::torch_float32() + } + } + + if (verbose) { + message("Building text encoder...") + } + # SD 2.1 uses the final-LN last hidden state (v-prediction path). The + # text encoder and VAE decoder compute in float32 (the denoise loop + # casts prompt embeds to unet_dtype and the latent to float32 before + # decode); only the UNet runs at unet_dtype. Up-casting a float16 + # artifact to float32 for these is lossless. + text_encoder <- text_encoder_native_from_safetensors( + file.path(diffusers_dir, "text_encoder"), apply_final_ln = TRUE, + verbose = FALSE) + text_encoder$to(device = torch::torch_device(devices$text_encoder), + dtype = torch::torch_float32()) + + if (verbose) { + message("Building UNet...") + } + unet <- unet_native_from_safetensors(file.path(diffusers_dir, "unet"), + verbose = FALSE) + unet$to(device = torch::torch_device(devices$unet), dtype = unet_dtype) + + if (verbose) { + message("Building VAE decoder...") + } + # decode = post_quant_conv + decoder (the SD VAE needs the 1x1 + # post_quant_conv the FLUX-derived native decoder omits) + decoder <- .sd_vae_decode_from_safetensors(file.path(diffusers_dir, "vae"), + latent_channels = 4L) + decoder$to(device = torch::torch_device(devices$decoder), + dtype = torch::torch_float32()) + + list(unet = unet, decoder = decoder, text_encoder = text_encoder) +} diff --git a/R/txt2img_sd21.R b/R/txt2img_sd21.R index 0241cd6..bca2733 100644 --- a/R/txt2img_sd21.R +++ b/R/txt2img_sd21.R @@ -24,6 +24,10 @@ #' Native text encoder has better GPU compatibility (especially Blackwell). #' @param use_native_unet Logical; if TRUE, uses native R torch UNet instead of TorchScript. #' Native UNet has better GPU compatibility (especially Blackwell). +#' @param diffusers_dir Optional path to a HuggingFace diffusers +#' directory (with `unet/`, `vae/`, `text_encoder/`). When set, the +#' pipeline is built natively from safetensors (no TorchScript), via +#' [sd_pipeline_from_safetensors()]. See [download_sd21()]. #' @param ... Additional parameters passed to the diffusion process. #' #' @return An image array and metadata @@ -42,7 +46,7 @@ txt2img_sd21 <- function(prompt, negative_prompt = NULL, img_dim = 768, filename = NULL, metadata_path = NULL, use_native_decoder = FALSE, use_native_text_encoder = FALSE, - use_native_unet = FALSE, ...) { + use_native_unet = FALSE, diffusers_dir = NULL, ...) { model_name <- "sd21" # Handle "auto" devices @@ -50,20 +54,43 @@ txt2img_sd21 <- function(prompt, negative_prompt = NULL, img_dim = 768, devices <- auto_devices(model_name) } - m2d <- models2devices(model_name = model_name, devices = devices, - unet_dtype_str = unet_dtype_str, - download_models = download_models) - devices <- m2d$devices - unet_dtype <- m2d$unet_dtype - device_cpu <- m2d$device_cpu - device_cuda <- m2d$device_cuda + device_cpu <- torch::torch_device("cpu") + device_cuda <- torch::torch_device("cuda") + if (!is.null(diffusers_dir)) { + # Native safetensors path: resolve devices/dtype with the pure + # helpers, skipping the .pt model verification models2devices runs. + # 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, + get_required_components(model_name)) + if (is.null(unet_dtype_str)) { + unet_dtype_str <- "float32" + } + unet_dtype <- setup_dtype(devices, unet_dtype_str) + } else { + m2d <- models2devices(model_name = model_name, devices = devices, + unet_dtype_str = unet_dtype_str, + download_models = download_models) + devices <- m2d$devices + unet_dtype <- m2d$unet_dtype + device_cpu <- m2d$device_cpu + device_cuda <- m2d$device_cuda + } if (is.null(pipeline)) { - pipeline <- load_pipeline(model_name = model_name, m2d = m2d, - unet_dtype_str = unet_dtype_str, - use_native_decoder = use_native_decoder, - use_native_text_encoder = use_native_text_encoder, - use_native_unet = use_native_unet) + if (!is.null(diffusers_dir)) { + # Native pipeline straight from HuggingFace diffusers + # safetensors (no TorchScript .pt); see download_sd21(). + pipeline <- sd_pipeline_from_safetensors(diffusers_dir, + model_name = model_name, devices = devices, + unet_dtype = unet_dtype) + } else { + pipeline <- load_pipeline(model_name = model_name, m2d = m2d, + unet_dtype_str = unet_dtype_str, + use_native_decoder = use_native_decoder, + use_native_text_encoder = use_native_text_encoder, + use_native_unet = use_native_unet) + } } # Start timing From 558b36a8f0b91061caf3d3367f1d6925a020e3cc Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 07:47:36 -0500 Subject: [PATCH 14/18] document(): man + NAMESPACE for the native SD 2.1 safetensors pipeline --- NAMESPACE | 2 ++ man/download_sd21.Rd | 20 ++++++++++++++++++ man/sd_pipeline_from_safetensors.Rd | 32 +++++++++++++++++++++++++++++ man/sd_pipeline_safetensors.Rd | 12 +++++++++++ man/txt2img_sd21.Rd | 7 ++++++- 5 files changed, 72 insertions(+), 1 deletion(-) create mode 100644 man/download_sd21.Rd create mode 100644 man/sd_pipeline_from_safetensors.Rd create mode 100644 man/sd_pipeline_safetensors.Rd diff --git a/NAMESPACE b/NAMESPACE index 9137c0e..e1920bc 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -13,6 +13,7 @@ export(download_flux1) export(download_flux2_klein) export(download_ltx2) export(download_model) +export(download_sd21) export(download_zimage_turbo) export(encode_bpe) export(encode_qwen) @@ -177,6 +178,7 @@ export(save_image) export(save_video) export(save_video_ltx23) export(scheduler_add_noise) +export(sd_pipeline_from_safetensors) export(sdxl_memory_profile) export(t5_encoder) export(text_encoder_native) diff --git a/man/download_sd21.Rd b/man/download_sd21.Rd new file mode 100644 index 0000000..f4288b6 --- /dev/null +++ b/man/download_sd21.Rd @@ -0,0 +1,20 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{download_sd21} +\alias{download_sd21} +\title{Download the Stable Diffusion 2.1 diffusers weights} +\usage{ +download_sd21(verbose = TRUE) +} +\arguments{ +\item{verbose}{Logical.} +} +\value{ +Invisibly, the diffusers directory (the parent of + \code{unet/}, \code{vae/}, \code{text_encoder/}). +} +\description{ +Fetches the UNet, VAE, and CLIP text encoder (config + safetensors) +from \code{stabilityai/stable-diffusion-2-1} into the hfhub cache. The +native tokenizer and DDIM scheduler need no downloads. About 5 GB +(fp32 UNet); one-time. +} diff --git a/man/sd_pipeline_from_safetensors.Rd b/man/sd_pipeline_from_safetensors.Rd new file mode 100644 index 0000000..5be6fe7 --- /dev/null +++ b/man/sd_pipeline_from_safetensors.Rd @@ -0,0 +1,32 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{sd_pipeline_from_safetensors} +\alias{sd_pipeline_from_safetensors} +\title{Assemble a native SD pipeline from a diffusers safetensors directory} +\usage{ +sd_pipeline_from_safetensors(diffusers_dir, model_name = "sd21", + devices = NULL, unet_dtype = NULL, verbose = TRUE) +} +\arguments{ +\item{diffusers_dir}{Directory with \code{unet/}, \code{vae/}, +\code{text_encoder/} subdirectories.} + +\item{model_name}{Currently "sd21" (SDXL pending its second encoder).} + +\item{devices}{Named list of component devices (\code{unet}, +\code{decoder}, \code{text_encoder}); defaults to all-CPU.} + +\item{unet_dtype}{A torch dtype for the UNet (default float16 on CUDA, +float32 on CPU).} + +\item{verbose}{Logical.} +} +\value{ +A list with \code{unet}, \code{decoder}, \code{text_encoder}. +} +\description{ +Builds the native UNet, VAE decoder, and CLIP text encoder from a +diffusers directory using the \code{*_from_safetensors} constructors, +places each on its component device, and returns the \code{$unet / +$decoder / $text_encoder} list the \code{txt2img_*} denoise loop +expects. +} diff --git a/man/sd_pipeline_safetensors.Rd b/man/sd_pipeline_safetensors.Rd new file mode 100644 index 0000000..4bf5bfc --- /dev/null +++ b/man/sd_pipeline_safetensors.Rd @@ -0,0 +1,12 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{sd_pipeline_safetensors} +\alias{sd_pipeline_safetensors} +\title{Native Stable Diffusion pipelines from diffusers safetensors} +\description{ +Assemble and run the native SD pipeline directly from a HuggingFace +diffusers directory (\code{unet/}, \code{vae/}, \code{text_encoder/}), +with no TorchScript \code{.pt} step - so it works on Blackwell and +loads the same weights everyone else uses. SD21 is wired end to end +here; SDXL still needs its second text encoder and added-conditioning +embeddings (tracked in tasks/todo.md). +} diff --git a/man/txt2img_sd21.Rd b/man/txt2img_sd21.Rd index 90362f1..4549222 100644 --- a/man/txt2img_sd21.Rd +++ b/man/txt2img_sd21.Rd @@ -9,7 +9,7 @@ txt2img_sd21(prompt, negative_prompt = NULL, img_dim = 768, pipeline = NULL, num_inference_steps = 50, guidance_scale = 7.5, seed = NULL, save_file = TRUE, filename = NULL, metadata_path = NULL, use_native_decoder = FALSE, use_native_text_encoder = FALSE, - use_native_unet = FALSE, ...) + use_native_unet = FALSE, diffusers_dir = NULL, ...) } \arguments{ \item{prompt}{A character string prompt describing the image to generate.} @@ -53,6 +53,11 @@ Native text encoder has better GPU compatibility (especially Blackwell).} \item{use_native_unet}{Logical; if TRUE, uses native R torch UNet instead of TorchScript. Native UNet has better GPU compatibility (especially Blackwell).} +\item{diffusers_dir}{Optional path to a HuggingFace diffusers +directory (with `unet/`, `vae/`, `text_encoder/`). When set, the +pipeline is built natively from safetensors (no TorchScript), via +[sd_pipeline_from_safetensors()]. See [download_sd21()].} + \item{...}{Additional parameters passed to the diffusion process.} } \value{ From 6b14e26706a8901f200428acd8d2553d3668a5ab Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 09:35:36 -0500 Subject: [PATCH 15/18] Fix native SD 2.1 UNet tiling: timestep embedding flip/shift The SD 2.1 native UNet computed its timestep embedding with flip_sin_to_cos=FALSE, downscale_freq_shift=1, which scrambles the sin/cos channel ordering the trained time_embedding weights expect. Standard diffusers SD (and the native SDXL UNet) use TRUE/0. The error was invisible to the constant-input parity test (GroupNorm sits at its bias and attention is uniform, so the spatial path is never exercised) but compounded through the network into spatially-tiled output - recognizable content repeated across the frame instead of one subject. Localized by comparing every native submodule to its callable TorchScript counterpart (all cos 1.0), then reconstructing the forward from TS submodules and sweeping wiring params: flip=TRUE,shift=0 -> cos 1.000000 vs the reference (was 0.82 at 64x64). Native UNet now matches TorchScript at cos 0.99999 across resolutions; SD 2.1 renders a coherent subject instead of tiles. test_unet.R gains a random-input parity check so this class of bug can't hide behind constant inputs again. --- NEWS.md | 13 ++++++++++--- R/unet.R | 11 ++++++++--- inst/tinytest/test_unet.R | 16 ++++++++++++++++ 3 files changed, 34 insertions(+), 6 deletions(-) diff --git a/NEWS.md b/NEWS.md index 65a522d..50e2424 100644 --- a/NEWS.md +++ b/NEWS.md @@ -40,9 +40,16 @@ `txt2img_sd21(diffusers_dir=)` uses it. The SD VAE decode now applies the `post_quant_conv` the FLUX-derived native decoder omitted (the decode was badly wrong without it). SD 2.1 defaults to float32 on this - path (fp16 attention overflows to NaN). A pre-existing spatial-coherence - ("tiling") bug in the native SD 2.1 UNet is tracked separately; it - affects the `.pt`-native path identically (same weights). + path (fp16 attention overflows to NaN). +* Fixed a native SD 2.1 UNet **tiling** bug (pre-existing; the + `.pt`-native path shared it). The timestep embedding used + `flip_sin_to_cos=FALSE, downscale_freq_shift=1`, scrambling the + sin/cos ordering the trained `time_embedding` weights expect (standard + diffusers SD, and the native SDXL UNet, use `TRUE/0`). Constant-input + parity tests could not see it (GroupNorm sits at its bias, attention is + uniform); it compounded through the spatial path into tiled output. + The native UNet now matches the TorchScript reference at cos 0.99999, + and `test_unet.R` gains a random-input parity check. All of the above is capability-**probed**, not version-pinned, so the fork requirement self-heals when the safetensors fixes reach CRAN diff --git a/R/unet.R b/R/unet.R index 6ed7cf2..5ed911c 100644 --- a/R/unet.R +++ b/R/unet.R @@ -258,10 +258,15 @@ unet_native <- torch::nn_module( # Get model dtype from weights model_dtype <- self$time_embedding_linear_1$weight$dtype - # Time embedding (computed in float32, then cast to model dtype) - # SD21 uses flip_sin_to_cos=FALSE, downscale_freq_shift=1 + # Time embedding (computed in float32, then cast to model dtype). + # SD 2.1 uses the standard diffusers timestep embedding + # (flip_sin_to_cos=TRUE, downscale_freq_shift=0), same as SDXL. The + # old FALSE/1 scrambled the sin/cos ordering the trained + # time_embedding weights expect - masked by constant inputs (so the + # parity test missed it) but compounding through the spatial path + # into tiled output. Verified cos 1.0 vs the TorchScript reference. t_emb <- timestep_embedding(timestep, self$block_out_channels[1], - flip_sin_to_cos = FALSE, downscale_freq_shift = 1L) + flip_sin_to_cos = TRUE, downscale_freq_shift = 0L) t_emb <- t_emb$to(dtype = model_dtype) t_emb <- self$time_embedding_linear_1(t_emb) t_emb <- torch::nnf_silu(t_emb) diff --git a/inst/tinytest/test_unet.R b/inst/tinytest/test_unet.R index 558feae..133b9b6 100644 --- a/inst/tinytest/test_unet.R +++ b/inst/tinytest/test_unet.R @@ -71,3 +71,19 @@ max_diff <- as.numeric(diff$max()) mean_diff <- as.numeric(diff$mean()) expect_true(max_diff < 0.1, info = paste("max_diff:", max_diff)) # deterministic inputs expect_true(mean_diff < 0.02, info = paste("mean_diff:", mean_diff)) + +# Random (non-constant) inputs at generation resolution. Constant inputs +# leave GroupNorm at its bias and attention uniform, so they cannot catch +# a scrambled timestep embedding; random inputs give cos ~0.82 with that +# bug and ~1.0 when correct. +torch::torch_manual_seed(123) +r_sample <- torch::torch_randn(c(1L, 4L, 64L, 64L)) +r_ctx <- torch::torch_randn(c(1L, 77L, 1024L)) +with_no_grad({ + ts_r <- ts_unet(r_sample, torch::torch_tensor(c(500L)), r_ctx) + nt_r <- native_unet$forward(r_sample, torch::torch_tensor(c(500L)), r_ctx) +}) +cos_sim <- as.numeric(torch::torch_sum(ts_r * nt_r)$item() / + (sqrt(torch::torch_sum(ts_r * ts_r)$item()) * + sqrt(torch::torch_sum(nt_r * nt_r)$item()))) +expect_true(cos_sim > 0.999, info = paste("random-input cos:", cos_sim)) From b6fda3737acd8f31283bfd1751ed75638887639b Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 09:48:57 -0500 Subject: [PATCH 16/18] Package the SD 2.1 .pt -> diffusers converter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit convert_sd21_pt_to_diffusers() rebuilds a diffusers-layout safetensors directory (unet/, vae/, text_encoder/ + derived CLIPTextConfig) from the cornball-ai/sd21-R TorchScript .pt files. A trace preserves the exact parameter tensors, so the output is bit-identical to the source at the chosen dtype (verified max|Δ|=0 vs the .pt). float16 keeps every file sub-2 GB (CRAN-safetensors readable). This is the provenance-clean build step for the hosted artifact now that upstream SD 2.1 was deprecated; cornball already hosts these OpenRAIL weights as .pt. --- R/convert_sd_pt.R | 109 +++++++++++++++++++++++++++++ inst/tinytest/test_convert_sd_pt.R | 32 +++++++++ 2 files changed, 141 insertions(+) create mode 100644 R/convert_sd_pt.R create mode 100644 inst/tinytest/test_convert_sd_pt.R diff --git a/R/convert_sd_pt.R b/R/convert_sd_pt.R new file mode 100644 index 0000000..96764cd --- /dev/null +++ b/R/convert_sd_pt.R @@ -0,0 +1,109 @@ +#' Convert cornball SD 2.1 TorchScript weights to a diffusers artifact +#' +#' Rebuilds a diffusers-layout directory (\code{unet/}, \code{vae/}, +#' \code{text_encoder/}) from the cornball-ai/sd21-R TorchScript +#' component \code{.pt} files, so the native safetensors pipeline +#' (\code{\link{download_sd21}} / \code{\link{sd_pipeline_from_safetensors}}) +#' can load SD 2.1 with no TorchScript. +#' +#' A TorchScript trace preserves the exact parameter tensors, so the +#' result is bit-identical to the source at the chosen dtype. This is the +#' provenance-clean way to build the hosted artifact: the upstream +#' \code{stabilityai/stable-diffusion-2-1} repo was deprecated, SD 2.1 is +#' CreativeML OpenRAIL++-M (redistributable), and cornball already hosts +#' these weights as \code{.pt}. At \code{float16} the components are all +#' sub-2 GB single files (unet ~1.7 GB, text_encoder ~0.65 GB, vae +#' ~0.16 GB), so they load on stock CRAN safetensors. +#' +#' @param pt_dir Directory holding \code{unet-cpu.pt}, +#' \code{decoder-cpu.pt}, \code{text_encoder-cpu.pt} (default: the +#' diffuseR \code{sd21} data location). +#' @param output_dir Output diffusers directory (default: the +#' \code{sd_pipeline_from_safetensors} / \code{download_sd21} location). +#' @param dtype \code{"float16"} (default, the hosted tier) or +#' \code{"float32"}. +#' @param verbose Logical. +#' +#' @return Invisibly, \code{output_dir}. +#' +#' @export +convert_sd21_pt_to_diffusers <- function(pt_dir = NULL, output_dir = NULL, + dtype = c("float16", "float32"), + verbose = TRUE) { + dtype <- match.arg(dtype) + if (!requireNamespace("safetensors", quietly = TRUE)) { + stop("The safetensors package is required to write the artifact.") + } + if (is.null(pt_dir)) { + pt_dir <- file.path(tools::R_user_dir("diffuseR", "data"), "sd21") + } + if (is.null(output_dir)) { + output_dir <- file.path(tools::R_user_dir("diffuseR", "data"), + "sd21-diffusers") + } + pt_dir <- path.expand(pt_dir) + td <- switch(dtype, float16 = torch::torch_float16(), + float32 = torch::torch_float32()) + + comps <- list( + list(pt = "unet-cpu.pt", strip = "^unet\\.", subdir = "unet", + file = "diffusion_pytorch_model.safetensors"), + list(pt = "decoder-cpu.pt", strip = "^vae\\.", subdir = "vae", + file = "diffusion_pytorch_model.safetensors"), + list(pt = "text_encoder-cpu.pt", strip = "^text_encoder\\.", + subdir = "text_encoder", file = "model.safetensors") + ) + absent <- Filter(function(c) !file.exists(file.path(pt_dir, c$pt)), comps) + if (length(absent)) { + stop("Missing TorchScript component(s) in ", pt_dir, ": ", + paste(vapply(absent, `[[`, character(1), "pt"), collapse = ", "), + ". Fetch them first (e.g. run a native sd21 pipeline once, or ", + "download_model(\"sd21\")).") + } + dir.create(output_dir, recursive = TRUE, showWarnings = FALSE) + + for (comp in comps) { + m <- torch::jit_load(file.path(pt_dir, comp$pt)) + params <- m$parameters + keys <- sub(comp$strip, "", names(params)) + sd <- stats::setNames( + lapply(params, function(p) p$detach()$to(dtype = td)$contiguous()), + keys) + d <- file.path(output_dir, comp$subdir) + dir.create(d, showWarnings = FALSE) + safetensors::safe_save_file(sd, file.path(d, comp$file)) + if (identical(comp$subdir, "text_encoder")) { + .sd21_write_clip_config(params, file.path(d, "config.json")) + } + if (verbose) { + message(sprintf(" %-13s %d tensors -> %s/%s", comp$subdir, + length(sd), comp$subdir, comp$file)) + } + rm(m, params, sd) + gc(verbose = FALSE) + } + + if (verbose) { + message("SD 2.1 diffusers artifact (", dtype, "): ", output_dir) + } + invisible(output_dir) +} + +# Derive a minimal CLIPTextConfig (what text_encoder_native_from_safetensors +# reads) from the text_encoder .pt parameters, keyed with the export prefix. +.sd21_write_clip_config <- function(params, path) { + g <- function(k) params[[paste0("text_encoder.", k)]] + tok <- g("text_model.embeddings.token_embedding.weight") + pos <- g("text_model.embeddings.position_embedding.weight") + fc1 <- g("text_model.encoder.layers.0.mlp.fc1.weight") + layers <- grep("encoder\\.layers\\.", names(params), value = TRUE) + n_layers <- length(unique(sub(".*layers\\.([0-9]+)\\..*", "\\1", layers))) + hidden <- as.integer(tok$shape[2]) + cfg <- list(vocab_size = as.integer(tok$shape[1]), hidden_size = hidden, + num_hidden_layers = as.integer(n_layers), + num_attention_heads = as.integer(hidden %/% 64L), + intermediate_size = as.integer(fc1$shape[1]), + max_position_embeddings = as.integer(pos$shape[1])) + jsonlite::write_json(cfg, path, auto_unbox = TRUE, pretty = TRUE) + invisible(cfg) +} diff --git a/inst/tinytest/test_convert_sd_pt.R b/inst/tinytest/test_convert_sd_pt.R new file mode 100644 index 0000000..1b97835 --- /dev/null +++ b/inst/tinytest/test_convert_sd_pt.R @@ -0,0 +1,32 @@ +# SD 2.1 .pt -> diffusers converter: CLIPTextConfig derivation from the +# text_encoder parameters, and the missing-component error. The full +# conversion (bit-identical to the source) is validated out of band +# against the cached .pt weights. + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +library(diffuseR) + +# --- config derivation from mock text_encoder params (small dims) ------------------ + +z <- function(...) torch::torch_zeros(c(...)) +p <- list() +p[["text_encoder.text_model.embeddings.token_embedding.weight"]] <- z(100L, 128L) +p[["text_encoder.text_model.embeddings.position_embedding.weight"]] <- z(77L, 128L) +for (i in 0:3) { + p[[sprintf("text_encoder.text_model.encoder.layers.%d.mlp.fc1.weight", i)]] <- + z(512L, 128L) +} +cfg <- diffuseR:::.sd21_write_clip_config(p, tempfile(fileext = ".json")) +expect_equal(cfg$vocab_size, 100L) +expect_equal(cfg$hidden_size, 128L) +expect_equal(cfg$max_position_embeddings, 77L) +expect_equal(cfg$num_hidden_layers, 4L) # distinct layer indices +expect_equal(cfg$num_attention_heads, 2L) # hidden / 64 +expect_equal(cfg$intermediate_size, 512L) + +# --- missing TorchScript components -> actionable error --------------------------- + +expect_error(convert_sd21_pt_to_diffusers(pt_dir = tempfile()), + pattern = "Missing") From 804b915e06e35b1c94238b94c67d9e47028d4ec1 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 09:48:57 -0500 Subject: [PATCH 17/18] document(): man + NAMESPACE for convert_sd21_pt_to_diffusers --- NAMESPACE | 1 + man/convert_sd21_pt_to_diffusers.Rd | 42 +++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 man/convert_sd21_pt_to_diffusers.Rd diff --git a/NAMESPACE b/NAMESPACE index e1920bc..76d0c37 100644 --- a/NAMESPACE +++ b/NAMESPACE @@ -5,6 +5,7 @@ export(bpe_tokenizer) export(clear_vram) export(clip_pooled_output) export(CLIPTokenizer) +export(convert_sd21_pt_to_diffusers) export(ddim_scheduler_create) export(ddim_scheduler_step) export(decode_bpe) diff --git a/man/convert_sd21_pt_to_diffusers.Rd b/man/convert_sd21_pt_to_diffusers.Rd new file mode 100644 index 0000000..e8fb0cf --- /dev/null +++ b/man/convert_sd21_pt_to_diffusers.Rd @@ -0,0 +1,42 @@ +% tinyrox says don't edit this manually, but it can't stop you! +\name{convert_sd21_pt_to_diffusers} +\alias{convert_sd21_pt_to_diffusers} +\title{Convert cornball SD 2.1 TorchScript weights to a diffusers artifact} +\usage{ +convert_sd21_pt_to_diffusers(pt_dir = NULL, output_dir = NULL, + dtype = c("float16", "float32"), verbose = TRUE) +} +\arguments{ +\item{pt_dir}{Directory holding \code{unet-cpu.pt}, +\code{decoder-cpu.pt}, \code{text_encoder-cpu.pt} (default: the +diffuseR \code{sd21} data location).} + +\item{output_dir}{Output diffusers directory (default: the +\code{sd_pipeline_from_safetensors} / \code{download_sd21} location).} + +\item{dtype}{\code{"float16"} (default, the hosted tier) or +\code{"float32"}.} + +\item{verbose}{Logical.} +} +\value{ +Invisibly, \code{output_dir}. +} +\description{ +Rebuilds a diffusers-layout directory (\code{unet/}, \code{vae/}, +\code{text_encoder/}) from the cornball-ai/sd21-R TorchScript +component \code{.pt} files, so the native safetensors pipeline +(\code{\link{download_sd21}} / \code{\link{sd_pipeline_from_safetensors}}) +can load SD 2.1 with no TorchScript. +} +\details{ +A TorchScript trace preserves the exact parameter tensors, so the +result is bit-identical to the source at the chosen dtype. This is the +provenance-clean way to build the hosted artifact: the upstream +\code{stabilityai/stable-diffusion-2-1} repo was deprecated, SD 2.1 is +CreativeML OpenRAIL++-M (redistributable), and cornball already hosts +these weights as \code{.pt}. At \code{float16} the components are all +sub-2 GB single files (unet ~1.7 GB, text_encoder ~0.65 GB, vae +~0.16 GB), so they load on stock CRAN safetensors. + +} From c6da0bbc9492404874a7d9aead9ea31121ca1312 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 10 Jul 2026 11:39:03 -0500 Subject: [PATCH 18/18] Point download_sd21 at the hosted cornball-ai/sd21-R artifact The upstream stabilityai/stable-diffusion-2-1 repo was deprecated; download_sd21 now fetches the fp16 diffusers safetensors from the cornball-ai/sd21-R dataset (diffusers/ prefix, repo_type="dataset"), built from our own .pt via convert_sd21_pt_to_diffusers and uploaded with hub_upload. No unet/vae config.json is needed (the native constructors use the SD 2.1 defaults); only the CLIP encoder config ships. --- R/sd_pipeline_safetensors.R | 55 ++++++++++++++----------------------- 1 file changed, 20 insertions(+), 35 deletions(-) diff --git a/R/sd_pipeline_safetensors.R b/R/sd_pipeline_safetensors.R index 374ace7..cba9c36 100644 --- a/R/sd_pipeline_safetensors.R +++ b/R/sd_pipeline_safetensors.R @@ -10,7 +10,7 @@ #' @name sd_pipeline_safetensors NULL -.sd21_repo <- "stabilityai/stable-diffusion-2-1" +.sd21_repo <- "cornball-ai/sd21-R" # The SD/SDXL AutoencoderKL applies a 1x1 post_quant_conv to the latent # before the decoder (decode(z) = decoder(post_quant_conv(z))). The @@ -48,19 +48,21 @@ sd_vae_decode <- torch::nn_module( m } -.sd21_files <- c("unet/config.json", - "unet/diffusion_pytorch_model.safetensors", - "vae/config.json", - "vae/diffusion_pytorch_model.safetensors", - "text_encoder/config.json", - "text_encoder/model.safetensors") +# Hosted on the cornball-ai/sd21-R dataset under diffusers/. fp16, sub-2 GB +# per file (CRAN-safetensors readable). No unet/vae config.json: the native +# constructors use the SD 2.1 defaults; only the CLIP encoder needs one. +.sd21_files <- c("diffusers/unet/diffusion_pytorch_model.safetensors", + "diffusers/vae/diffusion_pytorch_model.safetensors", + "diffusers/text_encoder/config.json", + "diffusers/text_encoder/model.safetensors") #' Download the Stable Diffusion 2.1 diffusers weights #' -#' Fetches the UNet, VAE, and CLIP text encoder (config + safetensors) -#' from \code{stabilityai/stable-diffusion-2-1} into the hfhub cache. The -#' native tokenizer and DDIM scheduler need no downloads. About 5 GB -#' (fp32 UNet); one-time. +#' Fetches the UNet, VAE, and CLIP text encoder from the +#' \code{cornball-ai/sd21-R} HuggingFace dataset (fp16 diffusers +#' safetensors, converted from the original OpenRAIL weights; the +#' upstream \code{stabilityai} repo was deprecated). About 2.5 GB, +#' one-time. The native tokenizer and DDIM scheduler need no downloads. #' #' @param verbose Logical. #' @@ -73,13 +75,13 @@ download_sd21 <- function(verbose = TRUE) { stop("The hfhub package is required to download model weights.") } have <- !is.null(tryCatch( - hfhub::hub_download(.sd21_repo, .sd21_files[[2]], - local_files_only = TRUE), + hfhub::hub_download(.sd21_repo, .sd21_files[[1]], + repo_type = "dataset", local_files_only = TRUE), error = function(e) NULL )) if (!have) { ok <- .ltx23_consent( - "Stable Diffusion 2.1 UNet + VAE + CLIP text encoder (~5 GB)" + "Stable Diffusion 2.1 UNet + VAE + CLIP text encoder (~2.5 GB)" ) if (!ok) { stop("Download cancelled.", call. = FALSE) @@ -88,27 +90,10 @@ download_sd21 <- function(verbose = TRUE) { message("Downloading Stable Diffusion 2.1 (diffusers safetensors)...") } } - dl <- function(f) { - tryCatch(hfhub::hub_download(.sd21_repo, f), error = function(e) { - msg <- conditionMessage(e) - if (grepl("404|not found|[Uu]nauthorized|[Ff]orbidden|401|403", - msg)) { - stop( - "Could not fetch ", .sd21_repo, " (", sub("\n.*", "", msg), - "). This repo is gated on HuggingFace; a 404 with a valid ", - "token means the license has not been accepted yet. To fix:\n", - " 1. Visit https://huggingface.co/", .sd21_repo, - " and accept the license.\n", - " 2. Ensure HF_TOKEN is set (Sys.getenv(\"HF_TOKEN\")).\n", - "Or pass an accessible diffusers directory straight to ", - "sd_pipeline_from_safetensors() / txt2img_sd21(diffusers_dir=).", - call. = FALSE) - } - stop(e) - }) - } - paths <- vapply(.sd21_files, dl, character(1)) - # the diffusers root is the parent of unet/ (paths[[1]] is unet/config.json) + paths <- vapply(.sd21_files, function(f) { + hfhub::hub_download(.sd21_repo, f, repo_type = "dataset") + }, character(1)) + # diffusers root = parent of unet/ (paths[[1]] is diffusers/unet/*.safetensors) diffusers_dir <- dirname(dirname(paths[[1]])) if (verbose) { message("SD 2.1 ready: ", diffusers_dir)