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")),
diff --git a/NAMESPACE b/NAMESPACE
index 1d44233..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)
@@ -13,6 +14,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)
@@ -80,6 +82,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)
@@ -169,14 +173,17 @@ export(preprocess_image)
export(quant_conv)
export(qwen_bpe_tokenizer)
export(qwen3_encoder)
+export(recommend)
export(save_frames)
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)
+export(text_encoder_native_from_safetensors)
export(text_encoder2_native)
export(tokenize_gemma3)
export(txt2img)
@@ -187,11 +194,14 @@ 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)
+export(vae_decoder_native_from_safetensors)
export(vocab_size)
export(vram_report)
export(write_wav)
diff --git a/NEWS.md b/NEWS.md
new file mode 100644
index 0000000..50e2424
--- /dev/null
+++ b/NEWS.md
@@ -0,0 +1,56 @@
+# 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.
+* 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).
+* `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).
+* 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
+(mlverse/safetensors#11, #13).
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/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/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 88233ba..ff9dad6 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)) {
@@ -221,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]]
}
}
@@ -235,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/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/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")) {
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/sd_pipeline_safetensors.R b/R/sd_pipeline_safetensors.R
new file mode 100644
index 0000000..cba9c36
--- /dev/null
+++ b/R/sd_pipeline_safetensors.R
@@ -0,0 +1,176 @@
+#' 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 <- "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
+# 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
+}
+
+# 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 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.
+#'
+#' @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[[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 (~2.5 GB)"
+ )
+ if (!ok) {
+ stop("Download cancelled.", call. = FALSE)
+ }
+ if (verbose) {
+ message("Downloading Stable Diffusion 2.1 (diffusers safetensors)...")
+ }
+ }
+ 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)
+ }
+ 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/st_caps.R b/R/st_caps.R
new file mode 100644
index 0000000..c4003a0
--- /dev/null
+++ b/R/st_caps.R
@@ -0,0 +1,183 @@
+#' 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, 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)",
+ 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)"))
+ lead <- if (fit) {
+ sprintf("%s is the best fit for your card but needs", precision)
+ } else {
+ sprintf("%s needs", precision)
+ }
+ sprintf(paste0(
+ "%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/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/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/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
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/R/unet_safetensors.R b/R/unet_safetensors.R
new file mode 100644
index 0000000..9667eb5
--- /dev/null
+++ b/R/unet_safetensors.R
@@ -0,0 +1,170 @@
+#' 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)
+}
+
+#' 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/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/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:
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")
diff --git a/inst/tinytest/test_recommend.R b/inst/tinytest/test_recommend.R
new file mode 100644
index 0000000..bf378f7
--- /dev/null
+++ b/inst/tinytest/test_recommend.R
@@ -0,0 +1,158 @@
+# 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)
+
+# --- 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)
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()
+}
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))
diff --git a/inst/tinytest/test_unet_safetensors.R b/inst/tinytest/test_unet_safetensors.R
new file mode 100644
index 0000000..79106c2
--- /dev/null
+++ b/inst/tinytest/test_unet_safetensors.R
@@ -0,0 +1,133 @@
+# 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)
+
+# --- 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()
+}
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.
+
+}
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/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/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/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/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/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/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.
+
+}
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/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{
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_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.
+
+}
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.
+}
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.
+}