Native safetensors recommender and SD 2.1 hosting pipeline

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

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)

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).

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

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

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

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

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

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,21 +224,26 @@ 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]]
         }
     }
 
     handle <- list(
                    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]])
     }
     )