From e9fab34126f31c45fcbf8e458c2bea5c40eb67da Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 10:04:50 -0500 Subject: [PATCH 1/9] resident_load(stream_text): stream Gemma 4 layers for small-card LTX-2.5 resident_load("ltx") already forwards resident_gb to ltx23_load_pipeline; stream_text passes stream to the Gemma 4 loader, so a serving entry can use the recommend("ltx25") placement on 8-12 GB cards. pinned_bytes now counts the streamed transformer blocks and encoder layers, which keep their staging outside the handle's set. --- R/resident.R | 30 ++++++++++++++++++++++-- inst/tinytest/test_resident_ltx_stream.R | 20 ++++++++++++++++ man/resident_load.Rd | 11 ++++++++- 3 files changed, 58 insertions(+), 3 deletions(-) create mode 100644 inst/tinytest/test_resident_ltx_stream.R diff --git a/R/resident.R b/R/resident.R index d46f6be..db9fad9 100644 --- a/R/resident.R +++ b/R/resident.R @@ -245,6 +245,13 @@ #' generate; the encoder rides to the card for its phase and back off, the #' way the pipeline's own components do, so it is never resident beside the #' transformer. +#' @param stream_text \code{ltx} with a Gemma 4 \code{text_encoder} only: +#' stream the encoder's layers to the card one at a time instead of +#' onloading it whole for the encode (\code{\link{load_gemma4_nf4}}'s +#' \code{stream}). \code{\link{recommend}("ltx25")} asks for it below +#' 12 GB. On small cards pass \code{resident_gb} too (through \code{...} +#' to \code{\link{ltx23_load_pipeline}}), so the transformer keeps that +#' many GB on the card and streams the rest of its blocks. #' #' @return A \code{diffuseR_resident} handle (an environment). Inspect it #' with \code{\link{resident_status}}; the fields of interest are the @@ -267,7 +274,8 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), device = "cuda", ..., verbose = TRUE, - text_encoder = NULL, tokenizer = NULL) { + text_encoder = NULL, tokenizer = NULL, + stream_text = FALSE) { model <- match.arg(model) ## NAMED ARGUMENTS RATHER THAN `...`, because `...` goes to the family ## loader and `ltx23_load_pipeline` has no `...` of its own -- an unknown @@ -279,6 +287,11 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", stop("text_encoder/tokenizer apply to the ltx family only; ", model, " loads its own", call. = FALSE) } + if (isTRUE(stream_text) && + (is.null(text_encoder) || !.is_gemma4_artifact(text_encoder))) { + stop("stream_text needs a Gemma 4 NF4 text_encoder artifact", + call. = FALSE) + } if (!torch::cuda_is_available()) { stop("resident_load() requires CUDA", call. = FALSE) } @@ -329,6 +342,13 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", res$gpu_components <- pipeline$gpu_components res$components <- names(.resident_components(pipeline)) res$pinned_bytes <- .resident_pinned_bytes(staging) + # Streamed LTX blocks (resident_gb) keep their pinned sets on the + # transformer, outside `staging` + block_stream <- pipeline$transformer$block_stream + if (!is.null(block_stream)) { + res$pinned_bytes <- res$pinned_bytes + + .resident_pinned_bytes(block_stream$staging) + } ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING. ## @@ -353,7 +373,8 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", if (.is_gemma4_artifact(text_encoder)) { if (verbose) message("Loading the Gemma 4 text encoder (pinned)...") res$text_encoder <- load_gemma4_text_encoder( - text_encoder, device = "cpu", pin = TRUE, verbose = verbose) + text_encoder, device = "cpu", pin = TRUE, + stream = isTRUE(stream_text), verbose = verbose) res$tokenizer <- gemma4_tokenizer(tokenizer) } else { if (verbose) message("Loading the Gemma3 text encoder (pinned)...") @@ -379,6 +400,11 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", res$pinned_bytes <- res$pinned_bytes + .resident_pinned_bytes(list(te_staging)) } + # A streamed encoder's layers carry one set each + if (!is.null(res$text_encoder$layer_stream)) { + res$pinned_bytes <- res$pinned_bytes + + .resident_pinned_bytes(res$text_encoder$layer_stream) + } } res$state <- "inactive" res$last_error <- NULL diff --git a/inst/tinytest/test_resident_ltx_stream.R b/inst/tinytest/test_resident_ltx_stream.R new file mode 100644 index 0000000..20e8797 --- /dev/null +++ b/inst/tinytest/test_resident_ltx_stream.R @@ -0,0 +1,20 @@ +# resident_load(stream_text = TRUE) streams a Gemma 4 encoder, so it is +# refused before any loading when there is no Gemma 4 artifact to stream. +# These checks run ahead of the CUDA check and need no GPU. + +library(tinytest) +library(diffuseR) + +expect_true("stream_text" %in% names(formals(resident_load))) +expect_false(formals(resident_load)$stream_text) + +expect_error(resident_load("ltx", stream_text = TRUE, + checkpoint_path = tempdir()), + "Gemma 4") +not_gemma4 <- file.path(tempdir(), "not-gemma4") +dir.create(not_gemma4, showWarnings = FALSE) +expect_error(resident_load("ltx", stream_text = TRUE, + checkpoint_path = tempdir(), + text_encoder = not_gemma4, tokenizer = not_gemma4), + "Gemma 4") +expect_error(resident_load("flux2", stream_text = TRUE), "Gemma 4") diff --git a/man/resident_load.Rd b/man/resident_load.Rd index ed2313f..4b8569f 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -9,7 +9,8 @@ resident_load( ..., verbose = TRUE, text_encoder = NULL, - tokenizer = NULL + tokenizer = NULL, + stream_text = FALSE ) } \arguments{ @@ -26,6 +27,14 @@ resident_load( \item{verbose}{Print progress messages.} +\item{stream_text}{\code{ltx} with a Gemma 4 \code{text_encoder} only: +stream the encoder's layers to the card one at a time instead of +onloading it whole for the encode (\code{\link{load_gemma4_nf4}}'s +\code{stream}). \code{\link{recommend}("ltx25")} asks for it below +12 GB. On small cards pass \code{resident_gb} too (through \code{...} +to \code{\link{ltx23_load_pipeline}}), so the transformer keeps that +many GB on the card and streams the rest of its blocks.} + \item{text_encoder,tokenizer}{\code{ltx} only: paths to the Gemma3 encoder artifact and the tokenizer directory, or for LTX-2.5 the Gemma 4 NF4 artifact (\code{\link{gemma4_quantize_nf4}}) for both, since it carries From 2c2cd066668eb79dacad61d1b18cdfd78b792c56 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 10:14:31 -0500 Subject: [PATCH 2/9] resident_load: count the streamed connector sets in pinned_bytes Below 7 GB resident the text connectors stream too, and their sets live on the connector modules, outside the handle's staging. Disk-staged pairs, which hold no host copy, count as zero. --- R/resident.R | 26 ++++++++++++++++++------ inst/tinytest/test_resident_ltx_stream.R | 20 ++++++++++++++++++ 2 files changed, 40 insertions(+), 6 deletions(-) diff --git a/R/resident.R b/R/resident.R index db9fad9..95cc529 100644 --- a/R/resident.R +++ b/R/resident.R @@ -115,6 +115,10 @@ total <- 0 for (st in staging) { for (pair in st) { + # Disk-staged pairs hold no host copy + if (is.null(pair$pinned)) { + next + } total <- total + prod(as.numeric(pair$pinned$shape)) * .dtype_bytes(pair$pinned$dtype) } @@ -122,6 +126,20 @@ total } +# Host bytes of an LTX pipeline's streamed sets, which live on the +# modules rather than in pipeline$staging: the transformer's streamed +# blocks (resident_gb) and, below 7 GB, the connectors' projections and +# blocks +.resident_stream_bytes <- function(pipeline) { + sets <- pipeline$transformer$block_stream$staging + con <- pipeline$connectors + if (!is.null(con)) { + sets <- c(sets, con$unit_stream, con$video_connector$layer_stream, + con$audio_connector$layer_stream) + } + .resident_pinned_bytes(sets) +} + # Bytes per element, keyed by the libtorch dtype name that # as.character() on a torch_dtype returns ("Float", "Half", "Byte", # "Long", ...), NOT the R constructor alias. Unknown dtypes fall back to @@ -342,12 +360,8 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", res$gpu_components <- pipeline$gpu_components res$components <- names(.resident_components(pipeline)) res$pinned_bytes <- .resident_pinned_bytes(staging) - # Streamed LTX blocks (resident_gb) keep their pinned sets on the - # transformer, outside `staging` - block_stream <- pipeline$transformer$block_stream - if (!is.null(block_stream)) { - res$pinned_bytes <- res$pinned_bytes + - .resident_pinned_bytes(block_stream$staging) + if (identical(model, "ltx")) { + res$pinned_bytes <- res$pinned_bytes + .resident_stream_bytes(pipeline) } ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING. diff --git a/inst/tinytest/test_resident_ltx_stream.R b/inst/tinytest/test_resident_ltx_stream.R index 20e8797..d19fcb9 100644 --- a/inst/tinytest/test_resident_ltx_stream.R +++ b/inst/tinytest/test_resident_ltx_stream.R @@ -18,3 +18,23 @@ expect_error(resident_load("ltx", stream_text = TRUE, text_encoder = not_gemma4, tokenizer = not_gemma4), "Gemma 4") expect_error(resident_load("flux2", stream_text = TRUE), "Gemma 4") + +# pinned_bytes counts the streamed sets that live on the modules +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +pair <- function(n) { + list(live = NULL, pinned = torch::torch_zeros(n, dtype = torch::torch_bfloat16())) +} +tr <- torch::nn_linear(2, 2) +tr$block_stream <- list(n_resident = 1L, staging = list(list(pair(10)), list(pair(6)))) +con <- torch::nn_module(initialize = function() { + self$video_connector <- torch::nn_linear(2, 2) + self$audio_connector <- torch::nn_linear(2, 2) +})() +con$unit_stream <- list(video_text_proj_in = list(pair(4))) +con$video_connector$layer_stream <- list(list(pair(3))) +con$audio_connector$layer_stream <- list(list(pair(2), list(live = NULL, pinned = NULL))) +bytes <- diffuseR:::.resident_stream_bytes(list(transformer = tr, connectors = con)) +expect_equal(bytes, 2 * (10 + 6 + 4 + 3 + 2)) +expect_equal(diffuseR:::.resident_stream_bytes(list(transformer = torch::nn_linear(2, 2))), 0) From 5730c7b4bcf1567b5b48e73ce2e7f1b9ab50fbc4 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 10:26:56 -0500 Subject: [PATCH 3/9] Bump version to 0.2.2.23 --- DESCRIPTION | 2 +- NEWS.md | 14 ++++++++++++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 75bab50..2756db3 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.22 +Version: 0.2.2.23 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), diff --git a/NEWS.md b/NEWS.md index b89cb82..5052d8f 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,17 @@ +# diffuseR 0.2.2.23 + +* `resident_load("ltx", stream_text = TRUE)` streams the Gemma 4 + encoder's layers, so a serving handle can take the small-card + placement: pass `resident_gb` too (forwarded to + `ltx23_load_pipeline()`). The handle's `pinned_bytes` now counts the + streamed transformer blocks, connector sets and encoder layers. + - Measured through `resident_load()`, `resident_activate()` and + `resident_generate()` with a 16 GB card held to a 12 GB board's free + memory: at `resident_gb = 5.38`, 768x512x121 renders in 74 s + (allocated peak 8.8 GiB). 960x960x145 needs `resident_gb` below 7.38 + on the GPU, and more than 37 GiB of host RAM in decode. The full + card, all resident: 768x512x121 in 81 s, allocated peak 13.6 GiB. + # diffuseR 0.2.2.22 * **FLUX 3 in less host RAM: fp8 Qwen3-VL.** `qwen3_quantize()` writes From 4bfea9662b897836fb346aebcb2104466c2776a7 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 14:24:30 -0500 Subject: [PATCH 4/9] LTX decode: copy the video to the R array a few frames at a time --- R/txt2vid_ltx23.R | 43 ++++++++++++++++++++------ inst/tinytest/test_ltx23_video_array.R | 21 +++++++++++++ 2 files changed, 54 insertions(+), 10 deletions(-) create mode 100644 inst/tinytest/test_ltx23_video_array.R diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 31daafe..9bf8339 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -59,6 +59,31 @@ ltx23_unpack_video_latents <- function(latents, num_frames, height, width, flatten(start_dim = 3L, end_dim = 4L) } +# Decoded video [B, 3, F, H, W] in [-1, 1] -> an R [F - skip, H, W, 3] +# double array in [0, 1], without the first `skip` frames. Copies +# `chunk` frames at a time into one preallocated array, so the host +# never holds a float32 copy of the whole clip beside the result (at +# 960x960x145: 1.6 GB float32, its contiguous copy, then 3.2 GB double) +.ltx23_video_array <- function(video, skip = 0L, chunk = 8L) { + f32 <- torch::torch_float32() + d <- as.integer(video$shape) + keep <- seq.int(skip + 1L, d[3]) + out <- array(0, dim = c(length(keep), d[4], d[5], 3L)) + for (s in seq.int(1L, length(keep), by = chunk)) { + idx <- s:min(length(keep), s + chunk - 1L) + x <- video[1L, , keep[idx[1]]:keep[idx[length(idx)]], , , + drop = FALSE]$squeeze(1L) + x <- ((x$to(dtype = f32) / 2 + 0.5)$clamp(0, 1))$ + permute(c(2L, 3L, 4L, 1L))$contiguous()$cpu() + out[idx, , , ] <- as.array(x) + # Each chunk's tensor and R copy are garbage now; collect them + # before the next, or they pile up to the size of the clip + rm(x) + gc(verbose = FALSE, full = FALSE) + } + out +} + # Pack audio latents [B, C, L, M] -> [B, L, C * M] ltx23_pack_audio_latents <- function(latents) { latents$transpose(2L, 3L)$flatten(start_dim = 3L, end_dim = 4L) @@ -1217,16 +1242,20 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, video_latents, vae$latents_mean, vae$latents_std ) video <- vae$decode(video_latents$to(dtype = vae_dtype)) - # [-1, 1] -> [0, 1], [B, 3, F, H, W] -> [F, H, W, 3] - video <- ((video$to(dtype = f32) / 2 + 0.5)$clamp(0, 1))[1,,,,] - video <- video$permute(c(2L, 3L, 4L, 1L))$cpu() }) if (verbose) { + if (grepl("^cuda", as.character(vae_device))) { + torch::cuda_synchronize() + } message(sprintf(" video decode: %.1fs", as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) phase_t0 <- Sys.time() } - result$video <- as.array(video) + # Head-free pixels (trim_frames is the conditioning overlap of a + # continuation); result$latents keeps the full sequence + result$video <- torch::with_no_grad( + .ltx23_video_array(video, skip = trim_frames) + ) rm(video, video_latents) offload("vae") gc(verbose = FALSE) @@ -1234,12 +1263,6 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, message(sprintf(" video to R array: %.1fs", as.numeric(difftime(Sys.time(), phase_t0, units = "secs")))) } - if (trim_frames > 0L) { - # Deliver head-free pixels (the conditioning overlap of a - # continuation); result$latents keeps the full sequence - result$video <- result$video[-seq_len(trim_frames),,,, - drop = FALSE] - } } if (decode_audio) { diff --git a/inst/tinytest/test_ltx23_video_array.R b/inst/tinytest/test_ltx23_video_array.R new file mode 100644 index 0000000..a45a05d --- /dev/null +++ b/inst/tinytest/test_ltx23_video_array.R @@ -0,0 +1,21 @@ +# .ltx23_video_array copies the decoded clip to an R array a few frames +# at a time; the result must match the whole-clip conversion exactly, +# for any chunk size and with leading frames trimmed + +library(tinytest) + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +torch::torch_manual_seed(1) +v <- torch::torch_randn(1, 3, 9, 6, 10)$mul(0.8)$to(dtype = torch::torch_bfloat16()) +whole <- as.array(((v$to(dtype = torch::torch_float32()) / 2 + 0.5)$ + clamp(0, 1))[1, , , , ]$permute(c(2L, 3L, 4L, 1L))$cpu()) + +for (ch in c(1L, 4L, 9L, 16L)) { + expect_identical(diffuseR:::.ltx23_video_array(v, chunk = ch), whole) + expect_identical(diffuseR:::.ltx23_video_array(v, skip = 3L, chunk = ch), + whole[-(1:3), , , , drop = FALSE]) +} +expect_identical(dim(diffuseR:::.ltx23_video_array(v, skip = 8L)), c(1L, 6L, 10L, 3L)) From 05baee48b1dfd7f5bd8df1b80a5a02cdffa4f04c Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 14:24:30 -0500 Subject: [PATCH 5/9] resident_load(stage_from = "disk"): stage the Gemma 4 encoder from disk too --- NEWS.md | 19 ++++++++++++++++++ R/resident.R | 25 ++++++++++++++++++++++-- inst/tinytest/test_resident_ltx_stream.R | 5 +++++ man/resident_load.Rd | 8 ++++++++ 4 files changed, 55 insertions(+), 2 deletions(-) diff --git a/NEWS.md b/NEWS.md index 5052d8f..cf45025 100644 --- a/NEWS.md +++ b/NEWS.md @@ -11,6 +11,25 @@ (allocated peak 8.8 GiB). 960x960x145 needs `resident_gb` below 7.38 on the GPU, and more than 37 GiB of host RAM in decode. The full card, all resident: 768x512x121 in 81 s, allocated peak 13.6 GiB. + - On an RTX 3060 12 GB, `resident_gb = 5.38` with + `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` (the native + allocator OOMed at 960 with 0.9 GiB reserved but unallocated): + 960x960x145 in 382 s (allocated peak 10.5 GiB, cgroup peak 48.8 + GiB), 768x512x121 in 140 s (7.9 GiB, 44.9 GiB). +* `resident_load("ltx", stage_from = "disk")` now stages the Gemma 4 + encoder from disk too, so the handle keeps no weights in host memory + (`pinned_bytes` 0). A Gemma 3 encoder is refused. +* LTX decode copies the video to the R array eight frames at a time + instead of converting the whole clip to float32 first: the host peak + above the result drops from about 5x to 1.3x its size, and the GPU + no longer holds a float32 copy of the clip. + - Both together, a 16 GB card held to the 3060's free 11692 MiB, + expandable segments, `resident_gb = 5.38`, `stream_text`, under a + 16 GB cgroup cap: 768x512x121 in 116 s (host RSS 1.1 GiB after + load, 3.4 GiB at the end; same output sd as host staging), + 960x960x145 in 235 s (RSS 5.3 GiB, allocated peak 9.2 GiB). Host + staging needed 29.7 GiB after load and was killed in decode at 960 + under a 37 GiB cap. # diffuseR 0.2.2.22 diff --git a/R/resident.R b/R/resident.R index 95cc529..5f6b36f 100644 --- a/R/resident.R +++ b/R/resident.R @@ -244,6 +244,14 @@ #' The pipeline is left \emph{inactive} (weights pinned on the host, no #' VRAM held). Call \code{\link{resident_activate}} before generating. #' +#' For \code{ltx}, \code{stage_from = "disk"} (through \code{...}) keeps +#' no weights in host memory: the pipeline and a Gemma 4 +#' \code{text_encoder} both read every onload from their memory-mapped +#' NF4 artifacts, and \code{pinned_bytes} counts only the small computed +#' buffers. It saves ~26 GB of host RAM for a re-read per phase (and per +#' step for the streamed blocks), which the OS page cache serves when it +#' can. A Gemma 3 \code{text_encoder} is refused. +#' #' @param model One of "flux1", "flux2", "zimage", "ltx", "sdxl". #' @param device Target CUDA device, e.g. "cuda" or "cuda:1". #' @param ... Passed to the family loader (\code{\link{flux_load_pipeline}}, @@ -310,6 +318,14 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", stop("stream_text needs a Gemma 4 NF4 text_encoder artifact", call. = FALSE) } + # stage_from goes to ltx23_load_pipeline through `...`; the Gemma 4 + # encoder loaded here follows it + stage_disk <- identical(list(...)$stage_from, "disk") + if (stage_disk && !is.null(text_encoder) && + !.is_gemma4_artifact(text_encoder)) { + stop("stage_from = \"disk\" with a text_encoder needs a Gemma 4 ", + "NF4 artifact", call. = FALSE) + } if (!torch::cuda_is_available()) { stop("resident_load() requires CUDA", call. = FALSE) } @@ -385,10 +401,15 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", # LTX-2.5 encodes with Gemma 4; its NF4 artifact carries the # tokenizer too. txt2vid_ltx2() picks the encode by tokenizer class. if (.is_gemma4_artifact(text_encoder)) { - if (verbose) message("Loading the Gemma 4 text encoder (pinned)...") + if (verbose) { + message("Loading the Gemma 4 text encoder (", + if (stage_disk) "from disk" else "pinned", ")...") + } res$text_encoder <- load_gemma4_text_encoder( text_encoder, device = "cpu", pin = TRUE, - stream = isTRUE(stream_text), verbose = verbose) + stream = isTRUE(stream_text), + stage_from = if (stage_disk) "disk" else "host", + verbose = verbose) res$tokenizer <- gemma4_tokenizer(tokenizer) } else { if (verbose) message("Loading the Gemma3 text encoder (pinned)...") diff --git a/inst/tinytest/test_resident_ltx_stream.R b/inst/tinytest/test_resident_ltx_stream.R index d19fcb9..197cf43 100644 --- a/inst/tinytest/test_resident_ltx_stream.R +++ b/inst/tinytest/test_resident_ltx_stream.R @@ -18,6 +18,11 @@ expect_error(resident_load("ltx", stream_text = TRUE, text_encoder = not_gemma4, tokenizer = not_gemma4), "Gemma 4") expect_error(resident_load("flux2", stream_text = TRUE), "Gemma 4") +# stage_from = "disk" covers the encoder too, which only Gemma 4 can do +expect_error(resident_load("ltx", stage_from = "disk", + checkpoint_path = tempdir(), + text_encoder = not_gemma4, tokenizer = not_gemma4), + "Gemma 4") # pinned_bytes counts the streamed sets that live on the modules if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 4b8569f..57085f2 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -75,6 +75,14 @@ first activation lands on GPU 1. Wrap the call in The pipeline is left \emph{inactive} (weights pinned on the host, no VRAM held). Call \code{\link{resident_activate}} before generating. +For \code{ltx}, \code{stage_from = "disk"} (through \code{...}) keeps +no weights in host memory: the pipeline and a Gemma 4 +\code{text_encoder} both read every onload from their memory-mapped +NF4 artifacts, and \code{pinned_bytes} counts only the small computed +buffers. It saves ~26 GB of host RAM for a re-read per phase (and per +step for the streamed blocks), which the OS page cache serves when it +can. A Gemma 3 \code{text_encoder} is refused. + } \examples{ \dontrun{ From dc15a36c07f774f8ea15c27f631f2ea2ae26b369 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 16:20:15 -0500 Subject: [PATCH 6/9] resident_load("flux3"), and FLUX 3 loader paths for the text encoder, VAE and tokenizer --- NEWS.md | 15 +++++++ R/resident.R | 44 +++++++++++++++--- R/txt2img_flux3.R | 69 +++++++++++++++++++++++++---- inst/tinytest/test_resident_flux3.R | 55 +++++++++++++++++++++++ man/flux3_load_pipeline.Rd | 15 +++++++ man/resident_load.Rd | 12 +++-- 6 files changed, 190 insertions(+), 20 deletions(-) create mode 100644 inst/tinytest/test_resident_flux3.R diff --git a/NEWS.md b/NEWS.md index cf45025..5801347 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,5 +1,20 @@ # diffuseR 0.2.2.23 +* `resident_load("flux3")` serves FLUX 3 Image Klein text-to-image + through `resident_activate()`, `resident_generate()` (to + `txt2img_flux3()`) and `resident_deactivate()`. `pinned_bytes` counts + the streamed Qwen3-VL layers and fp8 linears too. +* `flux3_load_pipeline(text_dir =, support_dir =)` read the Qwen3-VL + text encoder and the VAE and tokenizer from given directories + instead of the default artifact location and the HuggingFace cache, + so with `model_dir` every file it reads can come from an offline + mount. + - Measured through the resident handle with a 16 GB card held to an + RTX 3060's free 11692 MiB, expandable segments, `vram_gb = 12`, + `encoder = FALSE`, hub cache empty and offline: load 23 s, 1024x1024 + in 24.6 s twice (allocated peak 8.0 GiB), host RSS 14.9 GiB, + `pinned_bytes` 9.4 GiB. + * `resident_load("ltx", stream_text = TRUE)` streams the Gemma 4 encoder's layers, so a serving handle can take the small-card placement: pass `resident_gb` too (forwarded to diff --git a/R/resident.R b/R/resident.R index 5f6b36f..696d3b7 100644 --- a/R/resident.R +++ b/R/resident.R @@ -42,7 +42,8 @@ # Families that ship a pinned/staged loader. Keyed by the `model` name # used everywhere else in the package (see recommend()). -.resident_families <- c("flux1", "flux2", "zimage", "ltx", "sdxl", "sd21") +.resident_families <- c("flux1", "flux2", "flux3", "zimage", "ltx", "sdxl", + "sd21") #' Every nn_module field of a pipeline, by name #' @@ -140,6 +141,18 @@ .resident_pinned_bytes(sets) } +# Host bytes of a FLUX 3 pipeline's streamed sets, also off +# pipeline$staging: the Qwen3-VL layers (stream_text) and the fp8 +# linears past resident_gb, which stay on the host and copy in per +# forward +.resident_flux3_stream_bytes <- function(pipeline) { + total <- .resident_pinned_bytes(pipeline$text_encoder$layer_stream) + for (m in pipeline$fp8_split$streamed) { + total <- total + m$weight_fp8$numel() * .dtype_bytes(m$weight_fp8$dtype) + } + total +} + # Bytes per element, keyed by the libtorch dtype name that # as.character() on a torch_dtype returns ("Float", "Half", "Byte", # "Long", ...), NOT the R constructor alias. Unknown dtypes fall back to @@ -252,14 +265,18 @@ #' step for the streamed blocks), which the OS page cache serves when it #' can. A Gemma 3 \code{text_encoder} is refused. #' -#' @param model One of "flux1", "flux2", "zimage", "ltx", "sdxl". +#' @param model One of "flux1", "flux2", "flux3", "zimage", "ltx", "sdxl", +#' "sd21". #' @param device Target CUDA device, e.g. "cuda" or "cuda:1". #' @param ... Passed to the family loader (\code{\link{flux_load_pipeline}}, -#' \code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}}, +#' \code{\link{flux2_load_pipeline}}, \code{\link{flux3_load_pipeline}}, +#' \code{\link{zimage_load_pipeline}}, #' \code{\link{ltx23_load_pipeline}} or #' \code{\link{sdxl_load_pipeline}}). \code{ltx} requires #' \code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the -#' \code{\link{download_sdxl}} cache). +#' \code{\link{download_sdxl}} cache). \code{flux3} takes +#' \code{model_dir}, \code{text_dir} and \code{support_dir} to read +#' from given paths instead of the default locations. #' @param verbose Print progress messages. #' @param text_encoder,tokenizer \code{ltx} only: paths to the Gemma3 encoder #' artifact and the tokenizer directory, or for LTX-2.5 the Gemma 4 NF4 @@ -297,8 +314,8 @@ #' } #' #' @export -resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", - "sdxl", "sd21"), +resident_load <- function(model = c("flux2", "flux1", "flux3", "zimage", + "ltx", "sdxl", "sd21"), device = "cuda", ..., verbose = TRUE, text_encoder = NULL, tokenizer = NULL, stream_text = FALSE) { @@ -343,6 +360,7 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", loader <- switch(model, flux1 = flux_load_pipeline, flux2 = flux2_load_pipeline, + flux3 = flux3_load_pipeline, zimage = zimage_load_pipeline, ltx = ltx23_load_pipeline, sdxl = sdxl_load_pipeline, @@ -358,7 +376,14 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", } else { isTRUE(dots$phase_offload) } - pipeline <- loader(device = "cuda", verbose = verbose, ...) + if (identical(model, "flux3") && is.null(dots$pin)) { + # Pin in the loader, so its streamed sets (stream_text, + # resident_gb) are the pinned copies; .resident_pin would + # otherwise pin the modules again beside them + pipeline <- loader(device = "cuda", verbose = verbose, pin = TRUE, ...) + } else { + pipeline <- loader(device = "cuda", verbose = verbose, ...) + } staging <- .resident_pin(pipeline, verbose = verbose) # Pinning also evicted anything the loader had left resident, so the @@ -379,6 +404,10 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx", if (identical(model, "ltx")) { res$pinned_bytes <- res$pinned_bytes + .resident_stream_bytes(pipeline) } + if (identical(model, "flux3")) { + res$pinned_bytes <- res$pinned_bytes + + .resident_flux3_stream_bytes(pipeline) + } ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING. ## @@ -845,6 +874,7 @@ resident_generate <- function(res, prompt, ...) { gen <- switch(res$model, flux1 = txt2img_flux, flux2 = txt2img_flux2, + flux3 = txt2img_flux3, zimage = txt2img_zimage, ltx = txt2vid_ltx2, sdxl = txt2img_sdxl, diff --git a/R/txt2img_flux3.R b/R/txt2img_flux3.R index 8415193..6e16f91 100644 --- a/R/txt2img_flux3.R +++ b/R/txt2img_flux3.R @@ -44,6 +44,23 @@ NULL "bf16" } +# The precision of a Qwen3-VL text encoder directory: a quantized +# artifact's manifest format, else "bf16" for a text_encoder directory +# (config.json plus safetensors shards) +.flux3_text_dir_precision <- function(dir) { + for (p in c("fp8", "nf4")) { + if (.is_qwen3_artifact(dir, p)) { + return(p) + } + } + if (file.exists(file.path(dir, "config.json")) && + length(list.files(dir, pattern = "\\.safetensors$"))) { + return("bf16") + } + stop("Not a Qwen3-VL text encoder directory (no fp8/nf4 manifest, ", + "no config.json with safetensors): ", dir, call. = FALSE) +} + # Resolve a FLUX 3 file from the HuggingFace cache .flux3_cached <- function(file, revision = .flux3_revision) { if (!requireNamespace("hfhub", quietly = TRUE)) { @@ -105,6 +122,17 @@ NULL #' then NF4, artifact built by \code{\link{download_flux3_klein}}, #' else bf16), "fp8" (~4.4 GB of host memory instead of ~8 GB), #' "nf4" (~2.8 GB), or "bf16" (the cached source). +#' @param text_dir Qwen3-VL text encoder directory: an fp8 or NF4 +#' artifact (\code{\link{qwen3_quantize}}) or a bf16 +#' \code{text_encoder} directory. Its own format decides the +#' precision; a \code{text_precision} other than "auto" must match it. +#' NULL (default) finds it by \code{text_precision}. +#' @param support_dir Directory laid out like the upstream repository +#' holding at least \code{vae/config.json}, +#' \code{vae/diffusion_pytorch_model.safetensors} and +#' \code{text_encoder/tokenizer.json}, read instead of the +#' HuggingFace cache. With \code{model_dir} and \code{text_dir} it +#' names every file the pipeline reads, for offline mounts. #' @param verbose Logical. #' #' @return A \code{flux3_pipeline} list. @@ -118,8 +146,20 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", dtype = NULL, vram_gb = NULL, stream_text = NULL, resident_gb = NULL, text_precision = c("auto", "fp8", "nf4", "bf16"), + text_dir = NULL, support_dir = NULL, verbose = TRUE) { .flux2_rev(revision) + support <- function(file) { + if (is.null(support_dir)) { + return(.flux3_cached(file, revision)) + } + path <- file.path(path.expand(support_dir), file) + if (!file.exists(path)) { + stop("Missing ", file, " in support_dir ", support_dir, + call. = FALSE) + } + path + } precision <- match.arg(precision) if (precision != "bf16") { precision <- .flux_resolve_precision(precision, @@ -197,9 +237,8 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading FLUX 3 VAE...") } - vae_config <- jsonlite::fromJSON(.flux3_cached("vae/config.json", revision)) - vae_path <- .flux3_cached("vae/diffusion_pytorch_model.safetensors", - revision) + vae_config <- jsonlite::fromJSON(support("vae/config.json")) + vae_path <- support("vae/diffusion_pytorch_model.safetensors") latent_channels <- as.integer(vae_config$latent_channels %||% 32L) pipe$vae_bn_eps <- vae_config$batch_norm_eps %||% 1e-4 pipe$decoder <- load_flux2_vae_decoder(vae_path, @@ -222,13 +261,26 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", if (verbose) { message("Loading Qwen3-VL text encoder...") } - text_precision <- .flux3_resolve_text_precision(match.arg(text_precision)) - te_dir <- if (text_precision == "bf16") { + text_precision <- match.arg(text_precision) + if (!is.null(text_dir)) { + text_dir <- path.expand(text_dir) + found <- .flux3_text_dir_precision(text_dir) + if (!text_precision %in% c("auto", found)) { + stop("text_dir holds a ", found, " text encoder, not ", + text_precision, ": ", text_dir, call. = FALSE) + } + text_precision <- found + } else { + text_precision <- .flux3_resolve_text_precision(text_precision) + } + te_dir <- if (!is.null(text_dir)) { + text_dir + } else if (text_precision == "bf16") { for (f in c("text_encoder/model.safetensors.index.json", sprintf("text_encoder/model-%05d-of-00002.safetensors", 1:2))) { - .flux3_cached(f, revision) + support(f) } - dirname(.flux3_cached("text_encoder/config.json", revision)) + dirname(support("text_encoder/config.json")) } else { dir <- .flux3_text_artifact(text_precision) if (!.is_qwen3_artifact(dir, text_precision)) { @@ -244,8 +296,7 @@ flux3_load_pipeline <- function(model_dir = NULL, device = "cuda", dtype = if (text_device == "cpu") "float32" else "bfloat16", verbose = verbose ) - pipe$tokenizer <- qwen_bpe_tokenizer( - .flux3_cached("text_encoder/tokenizer.json", revision)) + pipe$tokenizer <- qwen_bpe_tokenizer(support("text_encoder/tokenizer.json")) components <- intersect(c("transformer", "decoder", "encoder"), names(pipe)) if (!identical(text_device, "cpu") && !stream_text) { diff --git a/inst/tinytest/test_resident_flux3.R b/inst/tinytest/test_resident_flux3.R new file mode 100644 index 0000000..5802671 --- /dev/null +++ b/inst/tinytest/test_resident_flux3.R @@ -0,0 +1,55 @@ +# resident_load("flux3") dispatches to flux3_load_pipeline and +# txt2img_flux3; the loader reads its text encoder, VAE and tokenizer +# from given paths (text_dir, support_dir), and pinned_bytes counts the +# sets FLUX 3 streams off pipeline$staging. No GPU or weights needed. + +library(tinytest) +library(diffuseR) + +expect_true("flux3" %in% diffuseR:::.resident_families) +expect_true("flux3" %in% eval(formals(resident_load)$model)) +expect_true(all(c("text_dir", "support_dir") %in% + names(formals(flux3_load_pipeline)))) +expect_error(resident_load("flux3", text_encoder = tempdir()), + "ltx family only") + +# text_dir precision comes from the directory itself +prec <- diffuseR:::.flux3_text_dir_precision +root <- tempfile("flux3-text-") +fp8 <- file.path(root, "fp8") +dir.create(fp8, recursive = TRUE) +writeLines('{"model": "qwen3", "format": "fp8", "shards": ["a.safetensors"]}', + file.path(fp8, "manifest.json")) +file.create(file.path(fp8, "a.safetensors")) +expect_equal(prec(fp8), "fp8") +bf16 <- file.path(root, "bf16") +dir.create(bf16) +writeLines("{}", file.path(bf16, "config.json")) +file.create(file.path(bf16, "model-00001-of-00002.safetensors")) +expect_equal(prec(bf16), "bf16") +empty <- file.path(root, "empty") +dir.create(empty) +expect_error(prec(empty), "Not a Qwen3-VL text encoder") +unlink(root, recursive = TRUE) + +# pinned_bytes: streamed Qwen3-VL layers plus the fp8 linears past +# resident_gb +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} +pair <- function(n) { + list(live = NULL, pinned = torch::torch_zeros(n, dtype = torch::torch_bfloat16())) +} +te <- torch::nn_linear(2, 2) +te$layer_stream <- list(list(pair(10)), list(pair(6), pair(4))) +lin <- function(n) { + m <- torch::nn_linear(2, 2) + m$weight_fp8 <- torch::torch_zeros(n, dtype = torch::torch_uint8()) + m +} +pipe <- list(text_encoder = te, + fp8_split = list(resident = list(lin(100)), + streamed = list(lin(7), lin(5)))) +expect_equal(diffuseR:::.resident_flux3_stream_bytes(pipe), 2 * 20 + 12) +expect_equal(diffuseR:::.resident_flux3_stream_bytes( + list(text_encoder = torch::nn_linear(2, 2))), 0) diff --git a/man/flux3_load_pipeline.Rd b/man/flux3_load_pipeline.Rd index 4fe8aa7..1d4f1b7 100644 --- a/man/flux3_load_pipeline.Rd +++ b/man/flux3_load_pipeline.Rd @@ -18,6 +18,8 @@ flux3_load_pipeline( stream_text = NULL, resident_gb = NULL, text_precision = c("auto", "fp8", "nf4", "bf16"), + text_dir = NULL, + support_dir = NULL, verbose = TRUE ) } @@ -71,6 +73,19 @@ then NF4, artifact built by \code{\link{download_flux3_klein}}, else bf16), "fp8" (~4.4 GB of host memory instead of ~8 GB), "nf4" (~2.8 GB), or "bf16" (the cached source).} +\item{text_dir}{Qwen3-VL text encoder directory: an fp8 or NF4 +artifact (\code{\link{qwen3_quantize}}) or a bf16 +\code{text_encoder} directory. Its own format decides the +precision; a \code{text_precision} other than "auto" must match it. +NULL (default) finds it by \code{text_precision}.} + +\item{support_dir}{Directory laid out like the upstream repository +holding at least \code{vae/config.json}, +\code{vae/diffusion_pytorch_model.safetensors} and +\code{text_encoder/tokenizer.json}, read instead of the +HuggingFace cache. With \code{model_dir} and \code{text_dir} it +names every file the pipeline reads, for offline mounts.} + \item{verbose}{Logical.} } \value{ diff --git a/man/resident_load.Rd b/man/resident_load.Rd index 57085f2..ceadb45 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -4,7 +4,7 @@ \title{Load a diffusion pipeline as a resident handle} \usage{ resident_load( - model = c("flux2", "flux1", "zimage", "ltx", "sdxl", "sd21"), + model = c("flux2", "flux1", "flux3", "zimage", "ltx", "sdxl", "sd21"), device = "cuda", ..., verbose = TRUE, @@ -14,16 +14,20 @@ resident_load( ) } \arguments{ -\item{model}{One of "flux1", "flux2", "zimage", "ltx", "sdxl".} +\item{model}{One of "flux1", "flux2", "flux3", "zimage", "ltx", "sdxl", +"sd21".} \item{device}{Target CUDA device, e.g. "cuda" or "cuda:1".} \item{...}{Passed to the family loader (\code{\link{flux_load_pipeline}}, -\code{\link{flux2_load_pipeline}}, \code{\link{zimage_load_pipeline}}, +\code{\link{flux2_load_pipeline}}, \code{\link{flux3_load_pipeline}}, +\code{\link{zimage_load_pipeline}}, \code{\link{ltx23_load_pipeline}} or \code{\link{sdxl_load_pipeline}}). \code{ltx} requires \code{checkpoint_path}; \code{sdxl} needs nothing (it defaults to the -\code{\link{download_sdxl}} cache).} +\code{\link{download_sdxl}} cache). \code{flux3} takes +\code{model_dir}, \code{text_dir} and \code{support_dir} to read +from given paths instead of the default locations.} \item{verbose}{Print progress messages.} From 7dbc7b3daf97060bb8657ae406282ac3b6618514 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 17:53:13 -0500 Subject: [PATCH 7/9] FLUX.2 klein and Z-Image in a 10.5 GiB GPU budget: streamed Qwen3, Z-Image resident_gb and per-block gc --- NEWS.md | 24 +++++++++++ R/dit_zimage.R | 13 ++++++ R/fp8_ltx23.R | 20 ++++++---- R/resident.R | 51 +++++++++++++++--------- R/txt2img_flux2.R | 18 ++++++++- R/txt2img_zimage.R | 48 +++++++++++++++++++--- inst/tinytest/test_resident_flux3.R | 10 ++++- inst/tinytest/test_resident_ltx_stream.R | 2 +- man/flux2_load_pipeline.Rd | 9 ++++- man/resident_load.Rd | 17 ++++---- man/zimage_load_pipeline.Rd | 14 ++++++- 11 files changed, 181 insertions(+), 45 deletions(-) diff --git a/NEWS.md b/NEWS.md index 5801347..c3fcb12 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,27 @@ +# diffuseR 0.2.2.24 + +* Image models in a ~10.5 GiB GPU budget (a 16 GB card that keeps a + speech pair resident beside them): + - `flux2_load_pipeline(stream_text = TRUE)` and + `zimage_load_pipeline(stream_text = TRUE)` stream the bf16 Qwen3 + layers through the encode, whose ~8 GB whole-encoder onload set + both models' GPU peak. `resident_load()` passes `stream_text` to + the flux2, flux3 and zimage loaders (it was LTX-only). + - `zimage_load_pipeline(resident_gb =)` keeps that many GB of fp8 + linears on the GPU and streams the rest per forward, as FLUX 3 does. + - Z-Image's transformer runs a minor gc after each block (the + `diffuseR.block_gc` option its loader already set): each block left + 2-3 GB of rotary and dequant temporaries until R collected them. + - A streamed fp8 linear copies its weight into a persistent per-shape + device buffer instead of a fresh tensor per call. + - Measured on a 16 GB card held to 10.5 GiB free, expandable + segments, 1024x1024, two renders each, same image every time: + FLUX.2 klein 9-10 s (allocated peak 6.2 GiB; it ran out of memory + in the encode before), Z-Image 16-18 s all resident (9.9 GiB), + 18 s at `resident_gb = 4` (8.3 GiB) or 3 (7.3 GiB). FLUX 3 on its + 8 GB plan, unchanged: 26-28 s, 4.5 GiB. LTX-2.5 960x960x145 at + `resident_gb = 4.38`, `stage_from = "disk"`: 242 s, 8.3 GiB. + # diffuseR 0.2.2.23 * `resident_load("flux3")` serves FLUX 3 Image Klein text-to-image diff --git a/R/dit_zimage.R b/R/dit_zimage.R index efbb82d..d18dd7b 100644 --- a/R/dit_zimage.R +++ b/R/dit_zimage.R @@ -119,9 +119,16 @@ zimage_transformer <- torch::nn_module( pad_rows <- self$x_pad_token$unsqueeze(1L)$expand(c(1L, img_pad, self$dim)) x_emb <- torch::torch_cat(list(x_emb, pad_rows), dim = 2L) } + # Each block leaves GBs of temporaries (rotary in float32, fp8 + # dequant) that only R's gc frees; collect per block as FLUX 3 does, + # or a tight card fills with garbage first + block_gc <- isTRUE(getOption("diffuseR.block_gc")) for (i in seq_len(length(self$noise_refiner))) { x_emb <- self$noise_refiner[[i]](x_emb, img_freqs, adaln_input = adaln, chunk_size = chunk_size) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } # Caption tokens: embed, pad, refine @@ -135,6 +142,9 @@ zimage_transformer <- torch::nn_module( for (i in seq_len(length(self$context_refiner))) { cap_emb <- self$context_refiner[[i]](cap_emb, cap_freqs, chunk_size = chunk_size) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } # Unified sequence, image first @@ -147,6 +157,9 @@ zimage_transformer <- torch::nn_module( unified <- self$layers[[i]](unified, unified_freqs, adaln_input = adaln, chunk_size = chunk_size) + if (block_gc) { + gc(verbose = FALSE, full = FALSE) + } } out <- self$all_final_layer[[patch_key]](unified, adaln) diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index fd9e3aa..120c8c8 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -68,14 +68,20 @@ ltx23_fp8_linear <- torch::nn_module( invisible(self) }, forward = function(x) { - # Transfer fp8 bytes first (no-op when resident), then cast into a - # persistent per-shape buffer and scale in place: zero fresh - # allocations per call, same math as the allocating upcast - w8 <- self$weight_fp8$to(device = x$device, non_blocking = TRUE) - w <- .ltx23_get_dequant_buffer( - c(self$out_features, self$in_features), x$dtype, x$device - ) + # A host (streamed) weight's fp8 bytes land in a persistent + # per-shape buffer too: a fresh $to() copy per call is device + # garbage until R collects it, which on a tight budget fills the + # card first. Then cast into the bf16 buffer and scale in place: + # zero fresh allocations per call, same math as the allocating upcast + shape <- c(self$out_features, self$in_features) + w8 <- self$weight_fp8 + w <- .ltx23_get_dequant_buffer(shape, x$dtype, x$device) torch::with_no_grad({ + if (w8$device$type != x$device$type) { + buf8 <- .ltx23_get_dequant_buffer(shape, w8$dtype, x$device) + buf8$copy_(w8, non_blocking = TRUE) + w8 <- buf8 + } w$copy_(w8) w$mul_(self$weight_scale$to(device = x$device, dtype = x$dtype)) }) diff --git a/R/resident.R b/R/resident.R index 696d3b7..b80b0b7 100644 --- a/R/resident.R +++ b/R/resident.R @@ -141,11 +141,11 @@ .resident_pinned_bytes(sets) } -# Host bytes of a FLUX 3 pipeline's streamed sets, also off -# pipeline$staging: the Qwen3-VL layers (stream_text) and the fp8 -# linears past resident_gb, which stay on the host and copy in per -# forward -.resident_flux3_stream_bytes <- function(pipeline) { +# Host bytes of a Qwen3-encoded pipeline's streamed sets (FLUX.2 klein, +# Z-Image, FLUX 3), also off pipeline$staging: the Qwen3 layers +# (stream_text) and FLUX 3's fp8 linears past resident_gb, which stay +# on the host and copy in per forward +.resident_qwen_stream_bytes <- function(pipeline) { total <- .resident_pinned_bytes(pipeline$text_encoder$layer_stream) for (m in pipeline$fp8_split$streamed) { total <- total + m$weight_fp8$numel() * .dtype_bytes(m$weight_fp8$dtype) @@ -288,13 +288,16 @@ #' generate; the encoder rides to the card for its phase and back off, the #' way the pipeline's own components do, so it is never resident beside the #' transformer. -#' @param stream_text \code{ltx} with a Gemma 4 \code{text_encoder} only: -#' stream the encoder's layers to the card one at a time instead of -#' onloading it whole for the encode (\code{\link{load_gemma4_nf4}}'s -#' \code{stream}). \code{\link{recommend}("ltx25")} asks for it below -#' 12 GB. On small cards pass \code{resident_gb} too (through \code{...} -#' to \code{\link{ltx23_load_pipeline}}), so the transformer keeps that -#' many GB on the card and streams the rest of its blocks. +#' @param stream_text Stream the text encoder's layers to the card one at +#' a time instead of onloading it whole for the encode. \code{ltx} with +#' a Gemma 4 \code{text_encoder} (\code{\link{load_gemma4_nf4}}'s +#' \code{stream}): \code{\link{recommend}("ltx25")} asks for it below +#' 12 GB; on small cards pass \code{resident_gb} too (through +#' \code{...} to \code{\link{ltx23_load_pipeline}}), so the transformer +#' keeps that many GB on the card and streams the rest of its blocks. +#' \code{flux2}, \code{flux3} and \code{zimage}: passed to the loader +#' as its \code{stream_text} when TRUE (FALSE leaves FLUX 3's own +#' \code{vram_gb} plan in charge). Refused for the other families. #' #' @return A \code{diffuseR_resident} handle (an environment). Inspect it #' with \code{\link{resident_status}}; the fields of interest are the @@ -330,7 +333,12 @@ resident_load <- function(model = c("flux2", "flux1", "flux3", "zimage", stop("text_encoder/tokenizer apply to the ltx family only; ", model, " loads its own", call. = FALSE) } - if (isTRUE(stream_text) && + qwen <- model %in% c("flux2", "flux3", "zimage") + if (isTRUE(stream_text) && !qwen && !identical(model, "ltx")) { + stop("stream_text applies to ltx, flux2, flux3 and zimage", + call. = FALSE) + } + if (isTRUE(stream_text) && identical(model, "ltx") && (is.null(text_encoder) || !.is_gemma4_artifact(text_encoder))) { stop("stream_text needs a Gemma 4 NF4 text_encoder artifact", call. = FALSE) @@ -376,14 +384,19 @@ resident_load <- function(model = c("flux2", "flux1", "flux3", "zimage", } else { isTRUE(dots$phase_offload) } - if (identical(model, "flux3") && is.null(dots$pin)) { + extra <- list() + if (qwen && is.null(dots$pin)) { # Pin in the loader, so its streamed sets (stream_text, # resident_gb) are the pinned copies; .resident_pin would # otherwise pin the modules again beside them - pipeline <- loader(device = "cuda", verbose = verbose, pin = TRUE, ...) - } else { - pipeline <- loader(device = "cuda", verbose = verbose, ...) + extra$pin <- TRUE + } + if (qwen && isTRUE(stream_text)) { + # Only when asked: FLUX 3 plans its own default from vram_gb + extra$stream_text <- TRUE } + pipeline <- do.call(loader, c(list(device = "cuda", verbose = verbose), + extra, dots)) staging <- .resident_pin(pipeline, verbose = verbose) # Pinning also evicted anything the loader had left resident, so the @@ -404,9 +417,9 @@ resident_load <- function(model = c("flux2", "flux1", "flux3", "zimage", if (identical(model, "ltx")) { res$pinned_bytes <- res$pinned_bytes + .resident_stream_bytes(pipeline) } - if (identical(model, "flux3")) { + if (model %in% c("flux2", "flux3", "zimage")) { res$pinned_bytes <- res$pinned_bytes + - .resident_flux3_stream_bytes(pipeline) + .resident_qwen_stream_bytes(pipeline) } ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING. diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R index f7306a9..3a38005 100644 --- a/R/txt2img_flux2.R +++ b/R/txt2img_flux2.R @@ -89,6 +89,11 @@ NULL #' @param encoder Logical. Also load the VAE encoder, which #' \code{\link{img2img_flux2}} needs (FALSE keeps a text-to-image #' pipeline's memory unchanged). +#' @param stream_text Logical. Stream the bf16 Qwen3 layers through the +#' GPU one at a time during the encode instead of onloading the ~8 GB +#' encoder whole, whose encode then dominates the GPU peak. For a +#' budget under ~11 GB (a 12 GB card, or a 16 GB card shared with +#' other models). Needs \code{phase_offload} on CUDA. #' #' @return A \code{flux2_pipeline} list. #' @@ -98,7 +103,7 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, pin = NULL, verbose = TRUE, revision = NULL, - encoder = FALSE) { + encoder = FALSE, stream_text = FALSE) { ## Checked before any device work, where a configuration mistake is ## cheapest to report. .flux2_rev(revision) @@ -213,15 +218,24 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", pipe$tokenizer <- qwen_bpe_tokenizer( .flux2_cached("tokenizer/tokenizer.json", revision)) + stream_text <- isTRUE(stream_text) && device == "cuda" && phase_offload && + !identical(text_device, "cpu") components <- c("transformer", "decoder") if (encoder) { components <- c(components, "encoder") } - if (!identical(text_device, "cpu")) { + if (!identical(text_device, "cpu") && !stream_text) { components <- c(components, "text_encoder") } pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, components, verbose = verbose) + if (stream_text) { + if (verbose) { + message("Qwen3 streams layer by layer") + } + pipe$staging <- c(pipe$staging, + list(text_encoder = .flux3_text_stream(pipe$text_encoder, pin))) + } structure(pipe, class = "flux2_pipeline") } diff --git a/R/txt2img_zimage.R b/R/txt2img_zimage.R index ce60fc5..9704ad8 100644 --- a/R/txt2img_zimage.R +++ b/R/txt2img_zimage.R @@ -61,6 +61,15 @@ NULL #' only that snapshot -- no \code{refs/} entry and no network. The #' transformer comes from \code{model_dir} and is unaffected. #' +#' @param stream_text Logical. Stream the bf16 Qwen3 layers through the +#' GPU one at a time during the encode instead of onloading the ~8 GB +#' encoder whole (see \code{\link{flux2_load_pipeline}}). Needs +#' \code{phase_offload} on CUDA. +#' @param resident_gb Numeric. GB of fp8 transformer weights kept on the +#' GPU through the denoise; the rest stay in (pinned) host memory and +#' copy in inside each forward (as \code{\link{flux3_load_pipeline}}). +#' \code{Inf} (default) keeps all ~5.9 GB resident. fp8 on CUDA only. +#' #' @return A \code{zimage_pipeline} list. #' #' @export @@ -68,7 +77,8 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4", "bf16"), text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, pin = NULL, - verbose = TRUE, revision = NULL) { + verbose = TRUE, revision = NULL, + stream_text = FALSE, resident_gb = Inf) { ## Checked before any device work, where a configuration mistake is ## cheapest to report. .flux2_rev(revision) @@ -100,8 +110,15 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", } if (device == "cuda") { # Footprint sized to the largest phase (the 8 GB Qwen3 encode), - # same reasoning as the FLUX.2 klein pipeline - .flux_gc_gates(footprint_gb = 12) + # same reasoning as the FLUX.2 klein pipeline. A split + # transformer runs in a smaller budget, where torch must collect + # well before the card is full (see flux3_load_pipeline) + if (is.finite(resident_gb)) { + .flux_gc_gates(footprint_gb = 12, allocated_rate = 0.8, + reserved_rate = 0.6) + } else { + .flux_gc_gates(footprint_gb = 12) + } } if (phase_offload) { @@ -132,6 +149,14 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", # Resident fp8 happens at onload time (the weights ride to the GPU # with the phase and back off after) pipe$fp8_resident <- identical(pipe$format, "fp8") && device == "cuda" + if (pipe$fp8_resident && is.finite(resident_gb)) { + pipe$fp8_split <- .flux3_fp8_split(pipe$transformer, resident_gb, pin) + if (verbose) { + message(sprintf("fp8 transformer: %.1f GB resident, %.1f GB streamed", + pipe$fp8_split$resident_gb, + pipe$fp8_split$streamed_gb)) + } + } if (verbose) { message("Loading VAE decoder...") @@ -176,12 +201,21 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", pipe$tokenizer <- qwen_bpe_tokenizer( .zimage_cached("tokenizer/tokenizer.json", revision)) + stream_text <- isTRUE(stream_text) && device == "cuda" && phase_offload && + !identical(text_device, "cpu") components <- c("transformer", "decoder") - if (!identical(text_device, "cpu")) { + if (!identical(text_device, "cpu") && !stream_text) { components <- c(components, "text_encoder") } pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, components, verbose = verbose) + if (stream_text) { + if (verbose) { + message("Qwen3 streams layer by layer") + } + pipe$staging <- c(pipe$staging, + list(text_encoder = .flux3_text_stream(pipe$text_encoder, pin))) + } structure(pipe, class = "zimage_pipeline") } @@ -380,7 +414,11 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, fp8_manual <- isTRUE(pipeline$fp8_resident) && is.null(staging[["transformer"]]) if (fp8_manual) { - .flux_fp8_to_device(transformer, device) + if (is.null(pipeline$fp8_split)) { + .flux_fp8_to_device(transformer, device) + } else { + .flux3_fp8_to_device(pipeline$fp8_split$resident, device) + } } if (verbose) { message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, diff --git a/inst/tinytest/test_resident_flux3.R b/inst/tinytest/test_resident_flux3.R index 5802671..95e2e51 100644 --- a/inst/tinytest/test_resident_flux3.R +++ b/inst/tinytest/test_resident_flux3.R @@ -50,6 +50,12 @@ lin <- function(n) { pipe <- list(text_encoder = te, fp8_split = list(resident = list(lin(100)), streamed = list(lin(7), lin(5)))) -expect_equal(diffuseR:::.resident_flux3_stream_bytes(pipe), 2 * 20 + 12) -expect_equal(diffuseR:::.resident_flux3_stream_bytes( +expect_equal(diffuseR:::.resident_qwen_stream_bytes(pipe), 2 * 20 + 12) +expect_equal(diffuseR:::.resident_qwen_stream_bytes( list(text_encoder = torch::nn_linear(2, 2))), 0) +# FLUX.2 klein and Z-Image stream the same Qwen3 layers (stream_text) +expect_true("stream_text" %in% names(formals(flux2_load_pipeline))) +expect_true("stream_text" %in% names(formals(zimage_load_pipeline))) +expect_false(formals(flux2_load_pipeline)$stream_text) +expect_false(formals(zimage_load_pipeline)$stream_text) +expect_equal(formals(zimage_load_pipeline)$resident_gb, quote(Inf)) diff --git a/inst/tinytest/test_resident_ltx_stream.R b/inst/tinytest/test_resident_ltx_stream.R index 197cf43..9f0ebb8 100644 --- a/inst/tinytest/test_resident_ltx_stream.R +++ b/inst/tinytest/test_resident_ltx_stream.R @@ -17,7 +17,7 @@ expect_error(resident_load("ltx", stream_text = TRUE, checkpoint_path = tempdir(), text_encoder = not_gemma4, tokenizer = not_gemma4), "Gemma 4") -expect_error(resident_load("flux2", stream_text = TRUE), "Gemma 4") +expect_error(resident_load("sdxl", stream_text = TRUE), "applies to ltx") # stage_from = "disk" covers the encoder too, which only Gemma 4 can do expect_error(resident_load("ltx", stage_from = "disk", checkpoint_path = tempdir(), diff --git a/man/flux2_load_pipeline.Rd b/man/flux2_load_pipeline.Rd index 6accf19..268ecbb 100644 --- a/man/flux2_load_pipeline.Rd +++ b/man/flux2_load_pipeline.Rd @@ -13,7 +13,8 @@ flux2_load_pipeline( pin = NULL, verbose = TRUE, revision = NULL, - encoder = FALSE + encoder = FALSE, + stream_text = FALSE ) } \arguments{ @@ -50,6 +51,12 @@ transformer comes from \code{model_dir} and is unaffected.} \item{encoder}{Logical. Also load the VAE encoder, which \code{\link{img2img_flux2}} needs (FALSE keeps a text-to-image pipeline's memory unchanged).} + +\item{stream_text}{Logical. Stream the bf16 Qwen3 layers through the +GPU one at a time during the encode instead of onloading the ~8 GB +encoder whole, whose encode then dominates the GPU peak. For a +budget under ~11 GB (a 12 GB card, or a 16 GB card shared with +other models). Needs \code{phase_offload} on CUDA.} } \value{ A \code{flux2_pipeline} list. diff --git a/man/resident_load.Rd b/man/resident_load.Rd index ceadb45..4ad5047 100644 --- a/man/resident_load.Rd +++ b/man/resident_load.Rd @@ -31,13 +31,16 @@ from given paths instead of the default locations.} \item{verbose}{Print progress messages.} -\item{stream_text}{\code{ltx} with a Gemma 4 \code{text_encoder} only: -stream the encoder's layers to the card one at a time instead of -onloading it whole for the encode (\code{\link{load_gemma4_nf4}}'s -\code{stream}). \code{\link{recommend}("ltx25")} asks for it below -12 GB. On small cards pass \code{resident_gb} too (through \code{...} -to \code{\link{ltx23_load_pipeline}}), so the transformer keeps that -many GB on the card and streams the rest of its blocks.} +\item{stream_text}{Stream the text encoder's layers to the card one at +a time instead of onloading it whole for the encode. \code{ltx} with +a Gemma 4 \code{text_encoder} (\code{\link{load_gemma4_nf4}}'s +\code{stream}): \code{\link{recommend}("ltx25")} asks for it below +12 GB; on small cards pass \code{resident_gb} too (through +\code{...} to \code{\link{ltx23_load_pipeline}}), so the transformer +keeps that many GB on the card and streams the rest of its blocks. +\code{flux2}, \code{flux3} and \code{zimage}: passed to the loader +as its \code{stream_text} when TRUE (FALSE leaves FLUX 3's own +\code{vram_gb} plan in charge). Refused for the other families.} \item{text_encoder,tokenizer}{\code{ltx} only: paths to the Gemma3 encoder artifact and the tokenizer directory, or for LTX-2.5 the Gemma 4 NF4 diff --git a/man/zimage_load_pipeline.Rd b/man/zimage_load_pipeline.Rd index 0139801..71b4fbc 100644 --- a/man/zimage_load_pipeline.Rd +++ b/man/zimage_load_pipeline.Rd @@ -12,7 +12,9 @@ zimage_load_pipeline( phase_offload = TRUE, pin = NULL, verbose = TRUE, - revision = NULL + revision = NULL, + stream_text = FALSE, + resident_gb = Inf ) } \arguments{ @@ -45,6 +47,16 @@ scheduler config, Qwen3 encoder, tokenizer) this pulls from the \code{snapshots//}, so the load works against a cache holding only that snapshot -- no \code{refs/} entry and no network. The transformer comes from \code{model_dir} and is unaffected.} + +\item{stream_text}{Logical. Stream the bf16 Qwen3 layers through the +GPU one at a time during the encode instead of onloading the ~8 GB +encoder whole (see \code{\link{flux2_load_pipeline}}). Needs +\code{phase_offload} on CUDA.} + +\item{resident_gb}{Numeric. GB of fp8 transformer weights kept on the +GPU through the denoise; the rest stay in (pinned) host memory and +copy in inside each forward (as \code{\link{flux3_load_pipeline}}). +\code{Inf} (default) keeps all ~5.9 GB resident. fp8 on CUDA only.} } \value{ A \code{zimage_pipeline} list. From 3d090be88bf04cdde5bf7a39f09ecde83f8d57e7 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 17:53:14 -0500 Subject: [PATCH 8/9] Bump version to 0.2.2.24 --- DESCRIPTION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/DESCRIPTION b/DESCRIPTION index 2756db3..0a973bf 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.2.23 +Version: 0.2.2.24 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), From 1fb2a8dca2458dd6f7095872a9450c29d553dcf9 Mon Sep 17 00:00:00 2001 From: TroyHernandez Date: Fri, 30 Oct 2026 17:59:45 -0500 Subject: [PATCH 9/9] NEWS: image-budget peaks with a 4000-byte prompt --- NEWS.md | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/NEWS.md b/NEWS.md index c3fcb12..6ccfa17 100644 --- a/NEWS.md +++ b/NEWS.md @@ -21,6 +21,11 @@ 18 s at `resident_gb = 4` (8.3 GiB) or 3 (7.3 GiB). FLUX 3 on its 8 GB plan, unchanged: 26-28 s, 4.5 GiB. LTX-2.5 960x960x145 at `resident_gb = 4.38`, `stage_from = "disk"`: 242 s, 8.3 GiB. + - The prompt sets the peak through the encode. With a 3990-byte + prompt: FLUX.2 klein unchanged (its encode stops at 512 tokens), + Z-Image at `resident_gb = 4` 8.8 GiB (also 512 tokens), FLUX 3 on + the 8 GB plan 6.3 GiB at 31 s; FLUX 3 on the 12 GB plan runs out + of memory in this budget. # diffuseR 0.2.2.23