Resident FLUX 3 and LTX streaming; image models in a 10.5 GiB GPU budget

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: diffuseR
 Title: Functional Interface to Diffusion Models in R
-Version: 0.2.2.22
+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")),

NEWS.md

@@ -1,3 +1,80 @@
+# 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.
+  - 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
+
+* `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
+  `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.
+  - 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
 
 * **FLUX 3 in less host RAM: fp8 Qwen3-VL.** `qwen3_quantize()` writes

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)

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

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
 #'
@@ -115,13 +116,43 @@
     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)
         }
     }
     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)
+}
+
+# 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)
+    }
+    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
@@ -226,14 +257,26 @@
 #' The pipeline is left \emph{inactive} (weights pinned on the host, no
 #' VRAM held). Call \code{\link{resident_activate}} before generating.
 #'
-#' @param model One of "flux1", "flux2", "zimage", "ltx", "sdxl".
+#' 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", "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
@@ -245,6 +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 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
@@ -264,10 +317,11 @@
 #' }
 #'
 #' @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) {
+                          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 +333,24 @@ 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)
     }
+    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)
+    }
+    # 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)
     }
@@ -296,6 +368,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,
@@ -311,7 +384,19 @@ resident_load <- function(model = c("flux2", "flux1", "zimage", "ltx",
     } else {
         isTRUE(dots$phase_offload)
     }
-    pipeline <- loader(device = "cuda", verbose = verbose, ...)
+    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
+        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
@@ -329,6 +414,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)
+    if (identical(model, "ltx")) {
+        res$pinned_bytes <- res$pinned_bytes + .resident_stream_bytes(pipeline)
+    }
+    if (model %in% c("flux2", "flux3", "zimage")) {
+        res$pinned_bytes <- res$pinned_bytes +
+            .resident_qwen_stream_bytes(pipeline)
+    }
 
     ## THE LTX TEXT ENCODER, LOADED ONCE AND KEPT OFF THE HANDLE'S STAGING.
     ##
@@ -351,9 +443,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, verbose = verbose)
+                text_encoder, device = "cpu", pin = TRUE,
+                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)...")
@@ -379,6 +477,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
@@ -784,6 +887,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,