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