diff --git a/DESCRIPTION b/DESCRIPTION index e7da035..73c3b11 100644 --- a/DESCRIPTION +++ b/DESCRIPTION @@ -1,6 +1,6 @@ Package: diffuseR Title: Functional Interface to Diffusion Models in R -Version: 0.2.0 +Version: 0.2.1.1 Authors@R: c( person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"), comment = c(ORCID = "0009-0005-4248-604X")), @@ -9,12 +9,13 @@ Authors@R: c( comment = "portions ported from the diffusers library (Apache-2.0); see inst/COPYRIGHTS"), person("Lightricks Ltd.", role = "cph", comment = "LTX checkpoint layout and pipeline constants; see inst/COPYRIGHTS")) -Description: A native R implementation of diffusion models providing a functional - interface to state-of-the-art generative AI. Inspired by Hugging Face's Python - 'diffusers' library, 'diffuseR' allows users to generate and manipulate images - using text prompts through models like Stable Diffusion without Python dependencies. - The package provides a streamlined, idiomatic R experience with support for multiple - diffusion schedulers and device acceleration. +Description: A native R implementation of diffusion models providing a + functional interface to state-of-the-art generative AI. Inspired by the + 'Python' library 'diffusers' from 'Hugging Face' + , 'diffuseR' generates and manipulates images + from text prompts using models such as 'Stable Diffusion', with no + 'Python' dependency. Supports multiple diffusion schedulers and device + acceleration. License: Apache License (>= 2) URL: https://github.com/cornball-ai/diffuseR BugReports: https://github.com/cornball-ai/diffuseR/issues diff --git a/NEWS.md b/NEWS.md index 5b5b8e4..0241eb0 100644 --- a/NEWS.md +++ b/NEWS.md @@ -1,3 +1,46 @@ +# diffuseR 0.2.1.1 + +Addresses the CRAN review of the 0.2.0 submission. + +* Every exported `.Rd` with a `\usage` block now documents its return + value: 50 `@return` tags added, chiefly to the `nn_module` generators + for the FLUX, FLUX.2, Z-Image, LTX-2.3 and Gemma3 ports. +* Examples: 14 of the 23 `\dontrun{}` blocks now run during check, and + were rewritten to be self-contained instead of referencing undefined + objects. The 9 that remain need model weights on disk and are + itemised in `cran-comments.md`. +* `ddim_scheduler_create()` was uncallable at its documented defaults: + `beta_schedule` was never passed through `match.arg()`, so `switch()` + errored on the length-3 default, and the `device` default was a + length-2 vector that `torch_tensor()` rejects. `ddim_scheduler_step()` + had the same missing `match.arg()` on `prediction_type`. Every + internal caller passed these explicitly, so the broken defaults went + unnoticed. `device` now defaults to `torch_device("cpu")`. +* `DESCRIPTION`: software names single-quoted ('Python', 'Stable + Diffusion', 'Hugging Face' with its URL) and the trailing whitespace + that had been folding into double spaces since the first commit + removed. +* `save_video()`'s mp4 example is no longer live: the encoder inherits + the session's stdin, which `R CMD check --as-cran` uses to feed the + example script to R, so it consumed part of the script. + +# diffuseR 0.2.0.1 + +* The FLUX-family image loaders (`flux_load_pipeline`, + `flux2_load_pipeline`, `zimage_load_pipeline`) now page-lock the + phase-swapped transformer, VAE decoder, and text encoder(s) at load, + so the per-generation CPU<->GPU moves run at DMA rate (offload becomes + a pointer swap). A new `pin` argument, `NULL` by default, resolves via + `options(diffuseR.pin_staging)` then the host-RAM-aware `recommend()` + decision. Resident-fp8 transformers (flux2/zimage) stage their fp8 + weight fields too. +* `flux_load_pipeline()` GPU-encodes T5-XXL (bfloat16) on 14 GB+ cards, + where its encode phase fits; smaller cards keep the float32 CPU + encode. `text_device` defaults to `NULL` (resolved from the VRAM + tier). An explicit `text_device = "cpu"` still encodes in place. +* Internal: the pinned-staging helpers lost their `ltx23` prefix + (`staging.R`); `recommend()` and the loaders share one `.pin_decision`. + # diffuseR 0.2.0 ## Serving diff --git a/R/audio_vae_ltx23.R b/R/audio_vae_ltx23.R index 1bb2901..6f4fc26 100644 --- a/R/audio_vae_ltx23.R +++ b/R/audio_vae_ltx23.R @@ -23,6 +23,10 @@ NULL #' @param stride Integer. #' @param causality_axis "height", "width", "width-compatibility", or "none". #' +#' @return Module whose forward(x) returns the convolved tensor, padded +#' so that each output frame depends only on current and earlier +#' input frames. +#' #' @export ltx23_audio_causal_conv2d <- torch::nn_module( "ltx23_audio_causal_conv2d", @@ -63,6 +67,9 @@ ltx23_audio_causal_conv2d <- torch::nn_module( #' @param in_channels,out_channels Integers. #' @param causality_axis Character. #' +#' @return Module whose forward(x) returns \code{x} plus the residual +#' branch, a tensor of the same shape as \code{x}. +#' #' @export ltx23_audio_resnet_block <- torch::nn_module( "ltx23_audio_resnet_block", @@ -109,6 +116,9 @@ ltx23_audio_resnet_block <- torch::nn_module( #' @param in_channels Integer. #' @param causality_axis Character. #' +#' @return Module whose forward(x) returns the tensor upsampled 2x by +#' nearest-neighbour interpolation and convolved. +#' #' @export ltx23_audio_upsample <- torch::nn_module( "ltx23_audio_upsample", @@ -139,6 +149,9 @@ ltx23_audio_upsample <- torch::nn_module( #' @param in_channels Integer. #' @param causality_axis Character. #' +#' @return Module whose forward(x) returns the strided convolution of +#' \code{x}, halving the downsampled axes. +#' #' @export ltx23_audio_downsample <- torch::nn_module( "ltx23_audio_downsample", @@ -168,6 +181,10 @@ ltx23_audio_downsample <- torch::nn_module( #' See \code{\link{ltx23_audio_decoder}}. #' @param in_channels Integer. Mel channels (2 = stereo). #' +#' @return Module whose forward(x) returns the encoded audio latent, a +#' tensor downsampled along time and mel axes with the configured +#' latent channel count. +#' #' @export ltx23_audio_encoder <- torch::nn_module( "ltx23_audio_encoder", @@ -269,6 +286,9 @@ ltx23_audio_encoder <- torch::nn_module( #' @param causality_axis Character. #' @param mel_bins Integer. Output mel bins (crop/pad target). #' +#' @return Module whose forward(x) returns the decoded mel +#' spectrogram reconstructed from an audio latent. +#' #' @export ltx23_audio_decoder <- torch::nn_module( "ltx23_audio_decoder", @@ -391,6 +411,11 @@ ltx23_audio_decoder <- torch::nn_module( #' See \code{\link{ltx23_audio_decoder}}. #' @param in_channels Integer. Mel input channels (2 = stereo). #' +#' @return Module bundling the audio encoder and decoder. Its +#' forward(z) is \code{decode(z)}, returning the mel spectrogram for a +#' latent; \code{$encode()} and \code{$decode()} are callable +#' separately. +#' #' @export ltx23_audio_vae <- torch::nn_module( "ltx23_audio_vae", diff --git a/R/auto_devices.R b/R/auto_devices.R index 80c5b7d..1d87ab8 100644 --- a/R/auto_devices.R +++ b/R/auto_devices.R @@ -25,16 +25,13 @@ #' @export #' #' @examples -#' \dontrun{ -#' # Auto-detect best configuration -#' devices <- auto_devices("sdxl") +#' # Force a strategy: no GPU or nvidia-smi needed. +#' str(auto_devices("sdxl", strategy = "cpu_only")) #' -#' # Use with models2devices -#' m2d <- models2devices("sdxl", devices = auto_devices("sdxl")) +#' str(auto_devices("sd21", strategy = "unet_gpu")) #' -#' # Force CPU-only -#' devices <- auto_devices("sdxl", strategy = "cpu_only") -#' } +#' # Auto-detect free VRAM and pick a strategy for this machine. +#' str(auto_devices("sdxl")) auto_devices <- function(model = "sdxl", strategy = "auto") { # Free-VRAM requirements in GB (float16 component sizes + overhead) requirements <- list(sd21 = list(full_gpu = 4, unet_gpu = 3), diff --git a/R/connectors_ltx23.R b/R/connectors_ltx23.R index 9435ed1..c9185b9 100644 --- a/R/connectors_ltx23.R +++ b/R/connectors_ltx23.R @@ -44,6 +44,10 @@ ltx23_per_token_rms_norm <- function(x, eps = 1e-6) { #' @param rope_type "split" (LTX-2.3) or "interleaved". #' @param num_attention_heads Integer. For the split per-head layout. #' +#' @return Module whose forward(batch_size, pos, device) returns +#' \code{list(cos_freqs, sin_freqs)}, the 1-D rotary tables for a +#' sequence of length \code{pos}. +#' #' @export ltx23_rotary_pos_embed_1d <- torch::nn_module( "ltx23_rotary_pos_embed_1d", @@ -142,6 +146,11 @@ ltx23_transformer_block_1d <- torch::nn_module( #' @param eps Numeric. Norm epsilon. #' @param gated_attention Logical. Per-head attention output gates. #' +#' @return Module whose forward(hidden_states, attention_mask, +#' attn_mask_binarize_threshold) returns +#' \code{list(hidden_states, attention_mask)}: the transformed +#' sequence and the (possibly binarized) mask that accompanies it. +#' #' @export ltx23_connector_transformer_1d <- torch::nn_module( "ltx23_connector_transformer_1d", @@ -266,6 +275,12 @@ ltx23_connector_transformer_1d <- torch::nn_module( #' (DiT inner dims: 4096 / 2048). #' @param proj_bias Logical. Projection bias (TRUE for LTX-2.3). #' +#' @return Module whose forward(text_encoder_hidden_states, +#' attention_mask) returns \code{list(video_text_embedding, +#' audio_text_embedding, attention_mask)}: the caption states adapted +#' for the video and audio cross-attention streams, plus the binary +#' mask to use with them. +#' #' @export ltx23_text_connectors <- torch::nn_module( "ltx23_text_connectors", diff --git a/R/ddim_scheduler.R b/R/ddim_scheduler.R index eda8a8c..056622a 100644 --- a/R/ddim_scheduler.R +++ b/R/ddim_scheduler.R @@ -44,14 +44,14 @@ #' \url{https://arxiv.org/abs/2010.02502} #' #' @examples -#' \dontrun{ -#' # Create a DDIM scheduler with custom parameters -#' scheduler <- ddim_scheduler_create( -#' num_train_timesteps = 1000, -#' num_inference_steps = 30, -#' eta = 0.5, -#' beta_schedule = "scaled_linear" -#' ) +#' if (torch::torch_is_installed()) { +#' scheduler <- ddim_scheduler_create( +#' num_train_timesteps = 1000, +#' num_inference_steps = 5, +#' eta = 0.5, +#' beta_schedule = "scaled_linear" +#' ) +#' scheduler$timesteps #' } #' @export ddim_scheduler_create <- function(num_train_timesteps = 1000, @@ -60,7 +60,8 @@ ddim_scheduler_create <- function(num_train_timesteps = 1000, beta_start = 0.00085, beta_end = 0.012, rescale_betas_zero_snr = FALSE, dtype = torch::torch_float32(), - device = c(torch::torch_device("cpu"), torch::torch_device("cuda"))) { + device = torch::torch_device("cpu")) { + beta_schedule <- match.arg(beta_schedule) betas <- switch(beta_schedule, "linear" = seq(beta_start, beta_end, length.out = num_train_timesteps), "scaled_linear" = seq(sqrt(beta_start), sqrt(beta_end), @@ -156,14 +157,18 @@ ddim_scheduler_create <- function(num_train_timesteps = 1000, #' \url{https://arxiv.org/abs/2202.00512} #' #' @examples -#' \dontrun{ -#' # Perform a denoising step -#' result <- ddim_scheduler_step( -#' model_output = model_output, -#' timestep = timestep, -#' sample = sample, -#' eta = 0, # Deterministic sampling -#' prediction_type = "epsilon") +#' if (torch::torch_is_installed()) { +#' scheduler <- ddim_scheduler_create(num_inference_steps = 5) +#' sample <- torch::torch_randn(c(1, 4, 8, 8)) +#' model_output <- torch::torch_randn(c(1, 4, 8, 8)) +#' result <- ddim_scheduler_step( +#' model_output = model_output, +#' timestep = scheduler$timesteps[1], +#' sample = sample, +#' schedule = scheduler, +#' eta = 0, # Deterministic sampling +#' prediction_type = "epsilon") +#' result$shape #' } #' @export ddim_scheduler_step <- function(model_output, timestep, sample, schedule, @@ -174,6 +179,7 @@ ddim_scheduler_step <- function(model_output, timestep, sample, schedule, prediction_type = c("epsilon", "sample", "v_prediction"), dtype = torch::torch_float32(), device = "cpu") { + prediction_type <- match.arg(prediction_type) # 1. get previous step value (= timestep + 1); i.e. python-indexing timestep_index <- torch::torch_tensor(timestep + 1, dtype = torch::torch_long(), device = torch::torch_device(device)) @@ -312,14 +318,16 @@ ddim_scheduler_step <- function(model_output, timestep, sample, schedule, #' specified timestep, with beta being the noise schedule. #' #' @examples -#' \dontrun{ -#' # Assuming we have latents, noise, and a scheduler -#' noised_latents <- scheduler_add_noise( -#' original_latents = latents, -#' noise = torch::torch_randn_like(latents), -#' timestep = scheduler$timesteps[1], -#' scheduler_obj = scheduler -#' ) +#' if (torch::torch_is_installed()) { +#' scheduler <- ddim_scheduler_create(num_inference_steps = 5) +#' latents <- torch::torch_randn(c(1, 4, 8, 8)) +#' noised_latents <- scheduler_add_noise( +#' original_latents = latents, +#' noise = torch::torch_randn_like(latents), +#' timestep = scheduler$timesteps[1], +#' scheduler_obj = scheduler +#' ) +#' noised_latents$shape #' } #' #' @export diff --git a/R/dit_flux2_modules.R b/R/dit_flux2_modules.R index 21460b9..8a7d5a6 100644 --- a/R/dit_flux2_modules.R +++ b/R/dit_flux2_modules.R @@ -26,6 +26,10 @@ NULL #' @param mod_param_sets Integer. Number of (shift, scale, gate) triples. #' @param bias Logical. #' +#' @return Module whose forward(temb) returns the modulation tensor +#' \code{linear(silu(temb))}, holding \code{mod_param_sets} triples of +#' (shift, scale, gate) along the last axis. +#' #' @export flux2_modulation <- torch::nn_module( "flux2_modulation", @@ -62,6 +66,9 @@ flux2_modulation <- torch::nn_module( #' @param mult Numeric. Inner dim multiplier (FLUX.2: 3.0). #' @param bias Logical. #' +#' @return Module whose forward(x) returns the SwiGLU-gated projection +#' of \code{x}, a tensor with the last axis of width \code{dim_out}. +#' #' @export flux2_feed_forward <- torch::nn_module( "flux2_feed_forward", @@ -94,6 +101,11 @@ flux2_feed_forward <- torch::nn_module( #' @param eps Numeric. RMS norm epsilon. #' @param bias Logical. #' +#' @return Module whose forward(hidden_states, image_rotary_emb, +#' chunk_size) returns the block output [B, S, query_dim]: attention +#' and MLP branches computed in parallel from one fused projection, +#' concatenated, and projected back by a second fused layer. +#' #' @export flux2_parallel_self_attention <- torch::nn_module( "flux2_parallel_self_attention", diff --git a/R/dit_flux_modules.R b/R/dit_flux_modules.R index afe04b0..3b83b5e 100644 --- a/R/dit_flux_modules.R +++ b/R/dit_flux_modules.R @@ -77,6 +77,10 @@ flux_ada_layer_norm_zero_single <- torch::nn_module( #' @param bias Logical. Bias on the projection (TRUE for FLUX.1, FALSE #' for FLUX.2). #' +#' @return Module whose forward(x, cond) returns \code{x} normalized and +#' then scaled and shifted by the conditioning embedding, a tensor of +#' the same shape as \code{x}. +#' #' @export flux_ada_layer_norm_continuous <- torch::nn_module( "flux_ada_layer_norm_continuous", @@ -112,6 +116,12 @@ flux_ada_layer_norm_continuous <- torch::nn_module( #' @param bias Logical. Bias on the linear projections (TRUE for FLUX.1, #' FALSE for FLUX.2). #' +#' @return Module whose forward(hidden_states, encoder_hidden_states, +#' image_rotary_emb, chunk_size) returns the attended image stream +#' [B, S, query_dim]. When \code{encoder_hidden_states} is supplied +#' (double-stream blocks) it returns \code{list(image, text)} instead, +#' each projected by its own output layer. +#' #' @export flux_attention <- torch::nn_module( "flux_attention", diff --git a/R/dit_ltx23.R b/R/dit_ltx23.R index 6ff9bc5..336e0c3 100644 --- a/R/dit_ltx23.R +++ b/R/dit_ltx23.R @@ -37,6 +37,11 @@ NULL #' @param gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn #' LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints). #' +#' @return Module whose forward(hidden_states, ...) returns +#' \code{list(sample, audio_sample)}: the predicted velocity for the +#' video latent tokens and, when the audio branch is active, for the +#' audio latent tokens (\code{audio_sample} is NULL otherwise). +#' #' @export ltx23_transformer <- torch::nn_module( "ltx23_transformer", diff --git a/R/dit_ltx23_modules.R b/R/dit_ltx23_modules.R index 0bf31ac..93288c6 100644 --- a/R/dit_ltx23_modules.R +++ b/R/dit_ltx23_modules.R @@ -17,6 +17,10 @@ NULL #' @param eps Numeric. Stability epsilon. #' @param elementwise_affine Logical. Learn a scale weight. #' +#' @return Module whose forward(x) returns \code{x} RMS-normalized over +#' the last axis and cast back to the input dtype, a tensor of the +#' same shape. +#' #' @export ltx23_rms_norm <- torch::nn_module( "ltx23_rms_norm", @@ -342,6 +346,9 @@ ltx23_ada_layer_norm_single <- torch::nn_module( #' @param rope_type "split" or "interleaved". #' @param apply_gated_attention Logical. Add per-head sigmoid output gates. #' +#' @return Module whose forward(hidden_states, ...) returns the attended +#' states [B, S, query_dim] after the output projection. +#' #' @export ltx23_attention <- torch::nn_module( "ltx23_attention", @@ -485,6 +492,10 @@ ltx23_attention <- torch::nn_module( #' @param dim Integer. Input/output dimension. #' @param mult Integer. Hidden dimension multiplier. #' +#' @return Module whose forward(x) returns the projected states passed +#' through a tanh-approximated GELU, a tensor of the same shape as +#' \code{x}. +#' #' @export ltx23_feed_forward <- torch::nn_module( "ltx23_feed_forward", @@ -544,6 +555,11 @@ ltx23_get_mod_params <- function(scale_shift_table, temb, batch_size) { #' @param rope_type "split" or "interleaved". #' @param perturbed_attn Logical. Enable the STG perturbation arguments. #' +#' @return Module whose forward(hidden_states, ...) returns +#' \code{list(hidden_states, audio_hidden_states)}, the video and audio +#' streams after self-attention, cross-attention and the feed-forward, +#' each the same shape as its input. +#' #' @export ltx23_transformer_block <- torch::nn_module( "ltx23_transformer_block", diff --git a/R/dit_zimage.R b/R/dit_zimage.R index 9776dab..efbb82d 100644 --- a/R/dit_zimage.R +++ b/R/dit_zimage.R @@ -31,6 +31,12 @@ #' @param patch_size Integer. Spatial patch size. Default 2. #' @param f_patch_size Integer. Temporal patch size. Default 1. #' +#' @return Module whose forward(x, t, cap_feats, chunk_size) returns the +#' predicted velocity for the single latent, a tensor [C, F, H, W] +#' matching the shape of \code{x}. Note that the checkpoint negates +#' this output and consumes a reversed timestep; see +#' \code{\link{txt2img_zimage}}. +#' #' @export zimage_transformer <- torch::nn_module( "zimage_transformer", diff --git a/R/dit_zimage_modules.R b/R/dit_zimage_modules.R index 64e3c35..b66fa64 100644 --- a/R/dit_zimage_modules.R +++ b/R/dit_zimage_modules.R @@ -26,6 +26,9 @@ NULL #' @param dim Integer. Model width. #' @param hidden_dim Integer. Hidden width. #' +#' @return Module whose forward(x) returns \code{w2(silu(w1(x)) * w3(x))}, +#' a tensor of the same shape as \code{x}. +#' #' @export zimage_feed_forward <- torch::nn_module( "zimage_feed_forward", @@ -52,6 +55,11 @@ zimage_feed_forward <- torch::nn_module( #' @param norm_eps Numeric. RMSNorm epsilon. Default 1e-5. #' @param modulation Logical. Whether the block is timestep-modulated. #' +#' @return Module whose forward(x, freqs, adaln_input, chunk_size) +#' returns the residual block output, a tensor of the same shape as +#' \code{x}. \code{adaln_input} is used only when the block was built +#' with \code{modulation = TRUE}. +#' #' @export zimage_block <- torch::nn_module( "zimage_block", @@ -109,6 +117,10 @@ zimage_block <- torch::nn_module( #' @param out_channels Integer. Patch output dim #' (patch^2 * f_patch * latent channels). #' +#' @return Module whose forward(x, c) returns the token-to-patch +#' projection [B, S, out_channels], ready for unpatchifying into a +#' latent. +#' #' @export zimage_final_layer <- torch::nn_module( "zimage_final_layer", @@ -138,6 +150,9 @@ zimage_final_layer <- torch::nn_module( #' @param mid_size Integer. Hidden width. The full model uses 1024. #' @param freq_size Integer. Sinusoid width. Default 256. #' +#' @return Module whose forward(t) returns the timestep embedding +#' [B, out_size]. +#' #' @export zimage_t_embedder <- torch::nn_module( "zimage_t_embedder", diff --git a/R/download_component.R b/R/download_component.R index 2f8b026..1d974f1 100644 --- a/R/download_component.R +++ b/R/download_component.R @@ -21,14 +21,15 @@ download_component <- function(model_name = "sd21", component, show_progress = TRUE) { filename <- paste0(component, "-", device, ".pt") if (!requireNamespace("hfhub", quietly = TRUE)) { - stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") + stop("Package 'hfhub' is required. Install it from CRAN before ", + "calling download_component() again.") } repo_id <- paste0("cornball-ai/", model_name, "-R") cached <- !overwrite && !is.null(tryCatch( - hfhub::hub_download(repo_id, filename, repo_type = "dataset", - local_files_only = TRUE), - error = function(e) NULL)) + hfhub::hub_download(repo_id, filename, repo_type = "dataset", + local_files_only = TRUE), + error = function(e) NULL)) if (!cached && !.ltx23_consent(paste0(model_name, " ", component, " (", filename, ")"))) { stop("Download cancelled.", call. = FALSE) diff --git a/R/download_flux.R b/R/download_flux.R index a22b18f..c76a441 100644 --- a/R/download_flux.R +++ b/R/download_flux.R @@ -103,7 +103,8 @@ download_flux1 <- function(quantize = TRUE, precision = c("nf4", "fp8"), local_files_only = TRUE), error = function(e) NULL ) - if (is.null(cached) || !.hub_all_cached(.flux1_repo, .flux1_transformer_files)) { + if (is.null(cached) || + !.hub_all_cached(.flux1_repo, .flux1_transformer_files)) { free <- .ltx23_disk_free_gb(path.expand("~")) if (!is.na(free) && free < 45) { warning(sprintf( diff --git a/R/download_flux2.R b/R/download_flux2.R index e1a1f5a..52e4a2b 100644 --- a/R/download_flux2.R +++ b/R/download_flux2.R @@ -75,7 +75,8 @@ download_flux2_klein <- function(quantize = TRUE, local_files_only = TRUE), error = function(e) NULL ) - if (is.null(cached) || !.hub_all_cached(.flux2_repo, .flux2_transformer_files)) { + if (is.null(cached) || + !.hub_all_cached(.flux2_repo, .flux2_transformer_files)) { free <- .ltx23_disk_free_gb(path.expand("~")) if (!is.na(free) && free < 25) { warning(sprintf( diff --git a/R/download_model.R b/R/download_model.R index 8134881..0c43f94 100644 --- a/R/download_model.R +++ b/R/download_model.R @@ -12,7 +12,8 @@ #' @keywords internal hf_download_pt <- function(model_name, filename, download = TRUE) { if (!requireNamespace("hfhub", quietly = TRUE)) { - stop("Package 'hfhub' is required. Install with: install.packages('hfhub')") + stop("Package 'hfhub' is required. Install it from CRAN before ", + "calling download_model() again.") } repo_id <- paste0("cornball-ai/", model_name, "-R") diff --git a/R/download_zimage.R b/R/download_zimage.R index 571189b..66e2c02 100644 --- a/R/download_zimage.R +++ b/R/download_zimage.R @@ -81,7 +81,8 @@ download_zimage_turbo <- function(quantize = TRUE, local_files_only = TRUE), error = function(e) NULL ) - if (is.null(cached) || !.hub_all_cached(.zimage_repo, .zimage_transformer_files)) { + if (is.null(cached) || + !.hub_all_cached(.zimage_repo, .zimage_transformer_files)) { free <- .ltx23_disk_free_gb(path.expand("~")) if (!is.na(free) && free < 35) { warning(sprintf( diff --git a/R/flowmatch_scheduler.R b/R/flowmatch_scheduler.R index a62f807..0b9bfd1 100644 --- a/R/flowmatch_scheduler.R +++ b/R/flowmatch_scheduler.R @@ -44,15 +44,15 @@ #' \url{https://arxiv.org/abs/2210.02747} #' #' @examples -#' \dontrun{ -#' # Create a FlowMatch scheduler -#' scheduler <- flowmatch_scheduler_create( -#' num_train_timesteps = 1000, -#' shift = 1.0 -#' ) +#' if (torch::torch_is_installed()) { +#' scheduler <- flowmatch_scheduler_create( +#' num_train_timesteps = 1000, +#' shift = 1.0 +#' ) #' -#' # Set timesteps for inference -#' scheduler <- flowmatch_set_timesteps(scheduler, num_inference_steps = 8) +#' # Set timesteps for inference +#' scheduler <- flowmatch_set_timesteps(scheduler, num_inference_steps = 8) +#' scheduler$timesteps #' } #' @export flowmatch_scheduler_create <- function(num_train_timesteps = 1000L, diff --git a/R/fp8_ltx23.R b/R/fp8_ltx23.R index ad07264..71a5fdb 100644 --- a/R/fp8_ltx23.R +++ b/R/fp8_ltx23.R @@ -39,6 +39,12 @@ ltx23_is_fp8_cast_key <- function(mapped_key) { #' @param out_features,in_features Integers. #' @param bias Logical. #' +#' @return Module whose forward(x) returns the linear projection of +#' \code{x}, with the fp8 weight bytes transferred and cast up to the +#' compute dtype for the matmul. Same result as an +#' \code{nn_linear} of the same shape, at a quarter of the resident +#' weight bytes. +#' #' @export ltx23_fp8_linear <- torch::nn_module( "ltx23_fp8_linear", @@ -54,7 +60,7 @@ ltx23_fp8_linear <- torch::nn_module( set_fp8_weight = function(weight, scale, pin = FALSE) { weight <- weight$to(device = "cpu") if (pin && torch::cuda_is_available()) { - weight <- .ltx23_pin_host(weight) + weight <- .pin_host(weight) } self$weight_fp8 <- weight self$weight_scale <- scale$to(device = "cpu", diff --git a/R/gemma3_text_encoder.R b/R/gemma3_text_encoder.R index 6cbc334..4ac4de5 100644 --- a/R/gemma3_text_encoder.R +++ b/R/gemma3_text_encoder.R @@ -405,6 +405,10 @@ gemma3_decoder_layer <- torch::nn_module( #' Full Gemma3 text encoder model. #' #' @param config Model configuration list. +#' @return Module whose forward(input_ids, ...) returns +#' \code{list(last_hidden_state, hidden_states)}: the final hidden +#' state [B, S, hidden_size] and the list of per-layer hidden states. +#' #' @export gemma3_text_model <- torch::nn_module( "Gemma3TextModel", @@ -687,7 +691,7 @@ load_gemma3_text_encoder <- function(model_path, device = "cpu", model$to(device = device) if (pin && device == "cpu") { - st <- .ltx23_pin_component(model) + st <- .pin_component(model) if (!is.null(st)) { attr(model, "staging") <- st if (verbose) { @@ -861,8 +865,8 @@ encode_with_gemma3 <- function(prompts, model = NULL, tokenizer = NULL, if (!is.null(staging) && device != "cpu") { cur <- tryCatch(staging[[1]]$live$device$type, error = function(e) NULL) if (!identical(cur, device)) { - .ltx23_staged_onload(staging, device) - on.exit(.ltx23_staged_offload(staging), add = TRUE) + .staged_onload(staging, device) + on.exit(.staged_offload(staging), add = TRUE) } } diff --git a/R/memory_flux.R b/R/memory_flux.R index 1b136e9..c8e3c74 100644 --- a/R/memory_flux.R +++ b/R/memory_flux.R @@ -1,8 +1,10 @@ #' FLUX Memory Profiles #' #' 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. +#' 12B transformer runs NF4 (~7 GB) or fp8 (~12 GB), phase-onloaded to +#' the GPU for denoise; the T5-XXL text encoder phase-onloads to the GPU +#' (bfloat16, pinned) on 14 GB+ cards and computes on the CPU (float32) +#' below that, where its ~9.8 GB encode phase does not fit. #' #' @name memory_flux NULL diff --git a/R/memory_sdxl.R b/R/memory_sdxl.R index 4b206c4..ca1cb62 100644 --- a/R/memory_sdxl.R +++ b/R/memory_sdxl.R @@ -25,13 +25,13 @@ #' @export #' #' @examples -#' \dontrun{ -#' # Auto-detect profile -#' profile <- sdxl_memory_profile() +#' # A stated VRAM budget is deterministic and needs no GPU. +#' str(sdxl_memory_profile(vram_gb = 8)) #' -#' # Specific VRAM -#' profile <- sdxl_memory_profile(vram_gb = 8) -#' } +#' str(sdxl_memory_profile(vram_gb = 24)) +#' +#' # Auto-detect free VRAM on this machine. +#' str(sdxl_memory_profile()) sdxl_memory_profile <- function(vram_gb = NULL) { # Auto-detect free VRAM if not provided if (is.null(vram_gb)) { diff --git a/R/nf4_ltx23.R b/R/nf4_ltx23.R index d2360d6..5d1e25f 100644 --- a/R/nf4_ltx23.R +++ b/R/nf4_ltx23.R @@ -202,6 +202,11 @@ ltx23_release_dequant_buffers <- function() { #' @param out_features,in_features Integers. #' @param bias Logical. #' +#' @return Module whose forward(x) returns the linear projection of +#' \code{x}, dequantizing the NF4 weight into a reusable buffer first. +#' Same result as an \code{nn_linear} of the same shape, at roughly an +#' eighth of the resident weight bytes. +#' #' @export ltx23_nf4_linear <- torch::nn_module( "ltx23_nf4_linear", diff --git a/R/quantize_flux.R b/R/quantize_flux.R index caa5238..948828d 100644 --- a/R/quantize_flux.R +++ b/R/quantize_flux.R @@ -167,6 +167,67 @@ NULL invisible(module) } +# Collect an fp8 transformer's plain-field weight tensors. weight_fp8 +# and weight_scale live outside parameters/buffers, so .pin_component's +# module walk misses them; this mirrors .flux_fp8_to_device's recursion +# and returns the tensors so they can be pinned and staged alongside the +# module. set_data (in .pin_component) mutates each in place, so the +# field reference stays valid and staged onload/offload covers them - +# replacing the per-generation .flux_fp8_to_device reassignment. +.flux_fp8_collect <- function(module, out = list()) { + for (name in names(module$children)) { + child <- module$children[[name]] + if (!is.null(child$weight_fp8)) { + out[[length(out) + 1L]] <- child$weight_fp8 + out[[length(out) + 1L]] <- child$weight_scale + } + out <- .flux_fp8_collect(child, out) + } + out +} + +# Pin the phase-swapped components of a flux-family pipeline so the +# per-generation CPU<->GPU moves run at DMA rate (see staging.R). +# Returns a named list of staging pairs keyed by component, or NULL when +# staging does not apply (opted out, no phase offload, or CPU device). +# Only the listed components are pinned; text encoders resident on the +# host (text_device == "cpu") are omitted by the caller since they +# compute in place. Resident-fp8 transformers (flux2/zimage) also pin +# their plain weight_fp8/weight_scale fields so staged transfer covers +# them; nf4 packed weights are buffers and ride along automatically. +.flux_build_staging <- function(pipe, pin, phase_offload, device, components, + verbose = TRUE) { + if (!(isTRUE(pin) && isTRUE(phase_offload) && + identical(device, "cuda") && torch::cuda_is_available())) { + return(NULL) + } + if (verbose) { + message("Pinning host staging buffers...") + } + staging <- list() + for (nm in components) { + module <- pipe[[nm]] + if (is.null(module)) { + next + } + extra <- if (identical(nm, "transformer") && + isTRUE(pipe$fp8_resident)) { + .flux_fp8_collect(module) + } else { + NULL + } + st <- .pin_component(module, extra = extra) + if (!is.null(st)) { + staging[[nm]] <- st + } + } + if (length(staging)) { + staging + } else { + NULL + } +} + # Family-specific hooks for quantization and loading .flux_family_hooks <- function(config) { family <- .flux_family(config) @@ -217,13 +278,15 @@ flux_quantize <- function(transformer_dir, output_dir = NULL, format <- match.arg(format) if (format == "fp8" && !.st_can_write("float8_e4m3fn")) { stop("The installed safetensors package cannot write float8 ", - "tensors (needs the float8 support pending in ", - "mlverse/safetensors#13; until it lands on CRAN, install ", - "remotes::install_github(\"cornball-ai/safetensors\") or ", - "use format = \"nf4\").", call. = FALSE) + "tensors (needs the float8 support from ", + "mlverse/safetensors#13, which is merged upstream but not ", + "yet on CRAN; install the development version of ", + "safetensors from GitHub, or use format = \"nf4\").", + call. = FALSE) } # Residents load into the compute dtype either way; bf16 halves the - # artifact but CRAN safetensors (<= 0.2.1) cannot write it yet + # artifact but the CRAN build of safetensors 0.2.1 cannot write it + # (the fix is merged upstream; the probe decides, not the version) resident_dtype <- if (.st_can_write("bfloat16")) { torch::torch_bfloat16() } else { @@ -380,9 +443,10 @@ flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16", format <- ckpt$format %||% "full" if (identical(format, "fp8") && !.st_can_write("float8_e4m3fn")) { stop("This fp8 artifact needs float8 support the installed ", - "safetensors lacks (pending in mlverse/safetensors#13). ", - "Install remotes::install_github(\"cornball-ai/safetensors\"), ", - "or rebuild the artifact as nf4.", call. = FALSE) + "safetensors lacks (mlverse/safetensors#13, merged upstream ", + "but not yet on CRAN). Install the development version of ", + "safetensors from GitHub, or rebuild the artifact as nf4.", + call. = FALSE) } hooks <- .flux_family_hooks(ckpt$config) diff --git a/R/quantize_gemma3.R b/R/quantize_gemma3.R index 0cf7e50..25f6f98 100644 --- a/R/quantize_gemma3.R +++ b/R/quantize_gemma3.R @@ -146,7 +146,7 @@ gemma3_quantize_nf4 <- function(model_path, output_dir = NULL, #' @param pin Logical. When loading to the CPU, page-lock the weights #' so \code{\link{encode_with_gemma3}} can swap the model to the GPU #' at DMA speed per encode and back for free (see -#' \code{\link{staging_ltx23}}). Default follows +#' \code{\link{staging}}). Default follows #' \code{options(diffuseR.pin_staging)}. #' @param verbose Logical. #' @@ -216,7 +216,7 @@ load_gemma3_nf4 <- function(artifact_dir, device = "cuda", model$to(device = device) model$eval() if (pin && device == "cpu") { - st <- .ltx23_pin_component(model) + st <- .pin_component(model) if (!is.null(st)) { attr(model, "staging") <- st if (verbose) { diff --git a/R/recommend.R b/R/recommend.R index 38a7045..c4e96ac 100644 --- a/R/recommend.R +++ b/R/recommend.R @@ -26,17 +26,17 @@ #' placement. #' #' The pinning decision: phase-swapped weights are page-locked host -#' copies (see \code{\link{staging_ltx23}}) that transfer at DMA rate - +#' copies (see \code{\link{staging}}) that transfer at DMA rate - #' but pinned pages are unswappable, so on small-RAM machines they turn #' memory pressure into OOM kills. \code{pin} is TRUE when available #' host RAM covers the model's pinned set twice over, FALSE below that, #' FALSE on the cpu tier (nothing stages), and TRUE when RAM cannot be -#' detected (page-locking already fails soft per component). Today the -#' LTX pipeline and Gemma3 encoder loaders take \code{pin} arguments -#' and stage pinned weights (see \code{\link{staging_ltx23}}); the -#' image-model loaders do not stage weights yet, so for them \code{pin} -#' is forward-looking policy with no consumer. -#' \code{options(diffuseR.pin_staging)} is the global switch. +#' detected (page-locking already fails soft per component). The LTX +#' pipeline, the Gemma3 encoder, and the FLUX-family image loaders +#' (flux1, flux2, zimage) take \code{pin} arguments and stage pinned +#' weights (see \code{\link{staging}}); the SD-family loaders place +#' components statically and do not phase-swap, so \code{pin} is inert +#' for them. \code{options(diffuseR.pin_staging)} is the global switch. #' #' @param model "sd21", "sdxl", "flux1", "flux2", "zimage", or "ltx". #' @param vram_gb Numeric or NULL. Free VRAM in GB; auto-detected via @@ -59,17 +59,20 @@ #' @export #' #' @examples -#' \dontrun{ -#' # Auto-detect VRAM and probe the installed safetensors -#' recommend("flux2") -#' -#' # A 16 GB card without float8 support: fp8 wanted, nf4 recommended +#' # Stating vram_gb and st_caps makes the policy deterministic: no GPU +#' # and no installed safetensors needed. #' r <- recommend("flux1", vram_gb = 16, #' st_caps = list(bfloat16 = TRUE, float8_e4m3fn = FALSE)) -#' r$precision # "nf4" +#' r$precision # "nf4": fp8 fits the card, but cannot be read #' r$fork_suggested # TRUE #' cat(r$note) # the fork-or-nf4 message -#' } +#' +#' # Same card, once safetensors can read float8 +#' recommend("flux1", vram_gb = 16, +#' st_caps = list(bfloat16 = TRUE, float8_e4m3fn = TRUE))$precision +#' +#' # Auto-detect VRAM and probe the installed safetensors +#' str(recommend("flux2")) recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", "ltx"), vram_gb = NULL, st_caps = NULL, host_ram_gb = NULL) { @@ -112,13 +115,8 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", host_ram_gb <- .detect_host_ram() } pinned_set <- .pinned_set_gb(model, chosen$precision) - pin <- if (isTRUE(chosen$cpu)) { - FALSE # cpu tier: nothing phase-stages, nothing to page-lock - } else if (is.na(host_ram_gb)) { - TRUE # undetectable: page-locking fails soft per component - } else { - host_ram_gb >= 2 * pinned_set - } + pin <- .pin_decision(model, chosen$precision, host_ram_gb, + cpu = isTRUE(chosen$cpu)) list( model = model, @@ -165,7 +163,7 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", sets <- list( sd21 = c(fp16 = 3), sdxl = c(fp16 = 8), - flux1 = c(bf16 = 43, fp8 = 31, nf4 = 26), # DiT + T5 fp32 host copy + flux1 = c(bf16 = 34, fp8 = 22, nf4 = 17), # DiT + T5 bf16 host copy flux2 = c(bf16 = 17, fp8 = 12, nf4 = 10), # DiT + Qwen3 bf16 zimage = c(bf16 = 21, fp8 = 14, nf4 = 12), # DiT + Qwen3 bf16 ltx = c(fp8 = 34, nf4 = 28) # checkpoint + Gemma3 NF4 @@ -178,6 +176,44 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", } } +# Host-RAM-aware pinning decision for a model at a resolved precision. +# TRUE when available host RAM covers the pinned set twice over; FALSE +# on the cpu tier (nothing stages, so nothing to page-lock); TRUE when +# RAM is undetectable (page-locking already fails soft per component). +# Shared by recommend() and the phase-offloading loaders. +.pin_decision <- function(model, precision, host_ram_gb = NULL, cpu = FALSE) { + if (isTRUE(cpu)) { + return(FALSE) + } + if (is.null(host_ram_gb)) { + host_ram_gb <- .detect_host_ram() + } + if (is.na(host_ram_gb)) { + return(TRUE) + } + host_ram_gb >= 2 * .pinned_set_gb(model, precision) +} + +# Resolve a loader's pin argument. Precedence: an explicit pin= wins, +# then options(diffuseR.pin_staging) (read with no default so an unset +# option falls through), then the RAM-aware .pin_decision. A "full" +# checkpoint is raw bf16, so it maps to the "bf16" pinned-set row. +.resolve_pin <- function(pin, model, format) { + if (!is.null(pin)) { + return(isTRUE(pin)) + } + opt <- getOption("diffuseR.pin_staging") + if (!is.null(opt)) { + return(isTRUE(opt)) + } + if (identical(format, "full")) { + precision <- "bf16" + } else { + precision <- format + } + .pin_decision(model, precision) +} + # 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. @@ -193,25 +229,38 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", # 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. +# +# text_device is the phase where the encoder computes: "cuda" onloads it +# to the GPU for the encode (fast, pinned), "cpu" keeps it resident on +# the host. flux1's T5-XXL needs ~9.8 GB for its encode phase, so it +# only GPU-encodes above gpu_encode_vram; the small flux2/zimage Qwen3 +# encodes on the GPU on every GPU tier (gpu_encode_vram = 0). .flux_family_tiers <- function(bf16_vram, fp8_vram, nf4_vram, nf4_tight_vram, max_hi, max_mid, max_lo, max_cpu, - attn_tight = NULL) { + attn_tight = NULL, gpu_encode_vram = 0) { + gpu_text <- function(min_vram) { + if (min_vram >= gpu_encode_vram) { + "cuda" + } else { + "cpu" + } + } list( list(precision = "bf16", min_vram = bf16_vram, needs = "bfloat16", devices = .dev_flux(TRUE), offload = TRUE, max_pixels = max_hi, - attn_chunk = NULL), + text_device = gpu_text(bf16_vram), 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), + text_device = gpu_text(fp8_vram), 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), + text_device = gpu_text(nf4_vram), 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), + text_device = gpu_text(nf4_tight_vram), 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) + text_device = "cpu", attn_chunk = NULL) ) } @@ -250,7 +299,7 @@ recommend <- function(model = c("sd21", "sdxl", "flux1", "flux2", "zimage", 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), + attn_tight = 2048L, gpu_encode_vram = 14), # 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, diff --git a/R/rope_ltx23.R b/R/rope_ltx23.R index 64b5e8a..fcc2098 100644 --- a/R/rope_ltx23.R +++ b/R/rope_ltx23.R @@ -111,6 +111,10 @@ ltx23_apply_split_rotary_emb <- function(x, freqs) { #' @param rope_type "split" (LTX 2.3) or "interleaved". #' @param num_attention_heads Integer. Needed for the split layout. #' +#' @return Module whose forward(coords, device) returns +#' \code{list(cos_freqs, sin_freqs)}, the two rotary tables to apply +#' to queries and keys. +#' #' @export ltx23_rotary_pos_embed <- torch::nn_module( "ltx23_rotary_pos_embed", diff --git a/R/save_image.R b/R/save_image.R index 457d857..79a8746 100644 --- a/R/save_image.R +++ b/R/save_image.R @@ -11,9 +11,11 @@ #' @export #' #' @examples -#' \dontrun{ -#' save_image(output_tensor, "sample.png") -#' } +#' img <- array(runif(32 * 32 * 3), dim = c(32, 32, 3)) +#' out <- file.path(tempdir(), "sample.png") +#' save_image(img, out) +#' file.exists(out) +#' unlink(out) save_image <- function(img, save_to = "output.png", normalize = TRUE) { # img_array <- tensor2image(img, normalize = normalize) dims <- dim(img) diff --git a/R/save_video.R b/R/save_video.R index 5456364..8b595b3 100644 --- a/R/save_video.R +++ b/R/save_video.R @@ -17,16 +17,21 @@ #' @export #' #' @examples -#' \dontrun{ -#' # Save as MP4 -#' save_video(video_array, "output.mp4", fps = 24) +#' video <- array(runif(4 * 16 * 16 * 3), dim = c(4, 16, 16, 3)) #' -#' # Save as GIF -#' save_video(video_array, "output.gif", fps = 10) +#' # Individual PNG frames need no external encoder. +#' frame_dir <- file.path(tempdir(), "frames") +#' save_video(video, frame_dir, format = "frames", verbose = FALSE) +#' length(list.files(frame_dir, pattern = "[.]png$")) +#' unlink(frame_dir, recursive = TRUE) #' -#' # Save as individual frames -#' save_video(video_array, "frames/", format = "frames") -#' } +#' # MP4 and GIF need an ffmpeg binary or the 'av' package. Not shown as +#' # a live example: both backends hand off to an ffmpeg process that +#' # inherits this session's stdin, and R CMD check feeds the example +#' # script to R on stdin, so the encoder eats a byte of the script and +#' # every later example parses one character short. +#' # save_video(video, "output.mp4", fps = 24) +#' # save_video(video, "output.gif", fps = 10) save_video <- function(video, file, fps = 24, format = NULL, backend = "auto", quality = 85, verbose = TRUE) { # Validate input diff --git a/R/serve.R b/R/serve.R index 1f0f334..ca19dea 100644 --- a/R/serve.R +++ b/R/serve.R @@ -75,13 +75,11 @@ #' #' @return Does not return normally; runs until interrupted. #' @export -serve <- function(port = 7812L, - model = c("flux2", "zimage", "flux1", "ltx"), - device = "cuda", token = NULL, - max_pixels = 1024L^2, max_frames = 161L, - max_steps = 50L, max_pixel_frames = NULL, - max_prompts = 32L, timeout = 300L, - max_body = 1024L^2, warmup = TRUE) { +serve <- function(port = 7812L, model = c("flux2", "zimage", "flux1", "ltx"), + device = "cuda", token = NULL, max_pixels = 1024L ^ 2, + max_frames = 161L, max_steps = 50L, + max_pixel_frames = NULL, max_prompts = 32L, timeout = 300L, + max_body = 1024L ^ 2, warmup = TRUE) { model <- match.arg(model) if (is.null(max_pixel_frames)) { # Joint video budget: full max_pixels only up to 121 frames; @@ -120,18 +118,16 @@ serve <- function(port = 7812L, srv <- serverSocket(port) on.exit(close(srv), add = TRUE) - message("diffuseR::serve listening on port ", port, - " (interrupt to stop)") + message("diffuseR::serve listening on port ", port, " (interrupt to stop)") repeat { con <- tryCatch( - socketAccept(srv, blocking = TRUE, open = "r+b", - timeout = timeout), + socketAccept(srv, blocking = TRUE, open = "r+b", timeout = timeout), error = function(e) { - message("accept error: ", conditionMessage(e)) - Sys.sleep(0.5) - NULL - } + message("accept error: ", conditionMessage(e)) + Sys.sleep(0.5) + NULL + } ) if (is.null(con)) { next @@ -143,23 +139,23 @@ serve <- function(port = 7812L, resp <- tryCatch( .dserve_route(req, state), error = function(e) { - r <- .dserve_err(500L, conditionMessage(e)) - if (grepl("CUDA out of memory", - conditionMessage(e), - fixed = TRUE)) { - attr(r, "fatal") <- TRUE - } - r - } + r <- .dserve_err(500L, conditionMessage(e)) + if (grepl("CUDA out of memory", + conditionMessage(e), + fixed = TRUE)) { + attr(r, "fatal") <- TRUE + } + r + } ) .dserve_send(con, resp$status, resp$content_type, resp$body) } }, - error = function(e) message("request error: ", conditionMessage(e)), - finally = { - try(close(con), silent = TRUE) - gc(verbose = FALSE) # bound dead handles; keep the warm pool - }) + error = function(e) message("request error: ", conditionMessage(e)), + finally = { + try(close(con), silent = TRUE) + gc(verbose = FALSE) # bound dead handles; keep the warm pool + }) if (exists("resp", inherits = FALSE) && isTRUE(attr(resp, "fatal"))) { message("CUDA out of memory: exiting for a clean supervisor restart") quit(save = "no", status = 70L) @@ -197,7 +193,11 @@ serve <- function(port = 7812L, } else { prec <- tryCatch(recommend("ltx")$precision, error = function(e) "nf4") - if (prec %in% names(avail)) avail[[prec]] else avail[[1]] + if (prec %in% names(avail)) { + avail[[prec]] + } else { + avail[[1]] + } } pipe <- ltx23_load_pipeline(pick, device = device, verbose = FALSE) te_dir <- file.path(data_dir, "gemma3-nf4") @@ -213,8 +213,7 @@ serve <- function(port = 7812L, } tok_dir <- dirname(hfhub::hub_download("Lightricks/LTX-2", "tokenizer/tokenizer.json", local_files_only = TRUE)) - te <- load_gemma3_text_encoder(te_dir, device = "cpu", - verbose = FALSE) + te <- load_gemma3_text_encoder(te_dir, device = "cpu", verbose = FALSE) tok <- gemma3_tokenizer(tok_dir) enc_dev <- if (identical(device, "cuda") && torch::cuda_is_available()) { @@ -237,9 +236,9 @@ serve <- function(port = 7812L, max_sequence_length = 1024L, device = enc_dev, verbose = FALSE) conn <- torch::with_no_grad(pipe$connectors( - raw$prompt_embeds$to(device = "cpu", - dtype = torch::torch_bfloat16()), - raw$prompt_attention_mask$to(device = "cpu"))) + raw$prompt_embeds$to(device = "cpu", + dtype = torch::torch_bfloat16()), + raw$prompt_attention_mask$to(device = "cpu"))) e <- list(video_text_embedding = conn$video_text_embedding, audio_text_embedding = conn$audio_text_embedding, attention_mask = conn$attention_mask) @@ -262,8 +261,7 @@ serve <- function(port = 7812L, pipe <- loader(device = device, verbose = FALSE) genfn <- switch(model, flux1 = txt2img_flux, flux2 = txt2img_flux2, zimage = txt2img_zimage) - generate <- function(prompt, width, height, seed = NULL, - steps = NULL) { + generate <- function(prompt, width, height, seed = NULL, steps = NULL) { args <- list(prompt = prompt, pipeline = pipe, width = as.integer(width), height = as.integer(height), seed = seed, @@ -279,8 +277,7 @@ serve <- function(port = 7812L, res } } - list(model = model, video = FALSE, generate = generate, - device = device) + list(model = model, video = FALSE, generate = generate, device = device) } } @@ -301,13 +298,15 @@ serve <- function(port = 7812L, } if (identical(req$method, "POST") && path == "/v1/images/generations") { if (isTRUE(state$video)) { - return(.dserve_err(400L, "this server hosts a video model; POST /v1/videos/generations")) + return(.dserve_err(400L, + "this server hosts a video model; POST /v1/videos/generations")) } return(.dserve_image(req, state)) } if (identical(req$method, "POST") && path == "/v1/videos/generations") { if (!isTRUE(state$video)) { - return(.dserve_err(400L, "this server hosts an image model; POST /v1/images/generations")) + return(.dserve_err(400L, + "this server hosts an image model; POST /v1/images/generations")) } return(.dserve_video(req, state)) } @@ -342,11 +341,11 @@ serve <- function(port = 7812L, .dserve_image <- function(req, state) { body <- .dserve_body(req) if (is.null(body) || !.dserve_prompt_ok(body$prompt)) { - return(.dserve_err(400L, "body must be JSON with a single string prompt")) + return(.dserve_err(400L, + "body must be JSON with a single string prompt")) } n <- .dserve_scalar(body$n, 1L) - if (is.na(suppressWarnings(as.integer(n))) || - as.integer(n) > 1L) { + if (is.na(suppressWarnings(as.integer(n))) || as.integer(n) > 1L) { return(.dserve_err(400L, "n > 1 is not supported")) } size <- .dserve_scalar(body$size, "1024x1024") @@ -359,15 +358,15 @@ serve <- function(port = 7812L, } if (anyNA(wh) || any(wh < 16L) || wh[1] * wh[2] > state$max_pixels) { return(.dserve_err(400L, sprintf( - "request exceeds limits (max %d pixels, min side 16)", - state$max_pixels))) + "request exceeds limits (max %d pixels, min side 16)", + state$max_pixels))) } steps <- .dserve_scalar(body$steps) if (!is.null(steps)) { steps <- suppressWarnings(as.integer(steps)) if (is.na(steps) || steps < 1L || steps > state$max_steps) { return(.dserve_err(400L, sprintf( - "steps must be between 1 and %d", state$max_steps))) + "steps must be between 1 and %d", state$max_steps))) } } seed <- .dserve_scalar(body$seed) @@ -381,15 +380,16 @@ serve <- function(port = 7812L, seed = seed, steps = steps) png <- png::writePNG(img) .dserve_json(list( - created = as.integer(Sys.time()), - data = list(list(b64_json = jsonlite::base64_enc(png))) - )) + created = as.integer(Sys.time()), + data = list(list(b64_json = jsonlite::base64_enc(png))) + )) } .dserve_video <- function(req, state) { body <- .dserve_body(req) if (is.null(body) || !.dserve_prompt_ok(body$prompt)) { - return(.dserve_err(400L, "body must be JSON with a single string prompt")) + return(.dserve_err(400L, + "body must be JSON with a single string prompt")) } w <- suppressWarnings(as.integer(.dserve_scalar(body$width, 768L))) h <- suppressWarnings(as.integer(.dserve_scalar(body$height, 512L))) @@ -399,8 +399,8 @@ serve <- function(port = 7812L, w * h > state$max_pixels || nf > state$max_frames || as.numeric(w) * h * nf > state$max_pixel_frames) { return(.dserve_err(400L, sprintf( - "request exceeds limits (max %d pixels, %d frames, %.0f pixel-frames)", - state$max_pixels, state$max_frames, state$max_pixel_frames))) + "request exceeds limits (max %d pixels, %d frames, %.0f pixel-frames)", + state$max_pixels, state$max_frames, state$max_pixel_frames))) } # frame_rate scales the audio-latent length inversely: bound it if (is.na(fr) || fr < 12 || fr > 60) { @@ -416,16 +416,16 @@ serve <- function(port = 7812L, out <- tempfile(fileext = ".mp4") on.exit(unlink(out), add = TRUE) txt2vid_ltx2( - prompt = body$prompt, - pipeline = state$pipe, - connector_embeds = state$embeds(body$prompt), - width = w, - height = h, - num_frames = nf, - frame_rate = fr, - seed = vseed, - device = state$device, dtype = "bfloat16", - filename = out, verbose = FALSE + prompt = body$prompt, + pipeline = state$pipe, + connector_embeds = state$embeds(body$prompt), + width = w, + height = h, + num_frames = nf, + frame_rate = fr, + seed = vseed, + device = state$device, dtype = "bfloat16", + filename = out, verbose = FALSE ) bytes <- readBin(out, "raw", n = file.size(out)) list(status = 200L, content_type = "video/mp4", body = bytes) @@ -496,8 +496,7 @@ serve <- function(port = 7812L, } reason <- switch(as.character(status), "200" = "OK", "400" = "Bad Request", "401" = "Unauthorized", - "404" = "Not Found", - "405" = "Method Not Allowed", + "404" = "Not Found", "405" = "Method Not Allowed", "413" = "Payload Too Large", "500" = "Internal Server Error", "Unknown") head <- paste0( diff --git a/R/st_caps.R b/R/st_caps.R index 97f2631..301b49e 100644 --- a/R/st_caps.R +++ b/R/st_caps.R @@ -1,9 +1,12 @@ #' 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: +#' The CRAN build of safetensors 0.2.1 reads bfloat16 but cannot write +#' it, and has no float8 support at all. Both fixes merged upstream on +#' 2026-07-31 (mlverse/safetensors#11 for bfloat16 write, #13 for +#' float8) without a version bump, so the installed version number +#' cannot tell you which build you have. That is why every gate here is +#' a runtime probe: write a tiny tensor, read it back, cache the answer. +#' Two capabilities matter and they differ: #' #' \itemize{ #' \item \emph{write} (\code{\link{flux_quantize}}'s internal @@ -111,11 +114,11 @@ NULL 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."), + "%s a safetensors newer than the one on CRAN: %s is ", + "merged upstream but not yet released. Install the ", + "development version from the mlverse/safetensors ", + "repository on GitHub, or press on with nf4: same ", + "weights, slightly lower precision, and it just works."), lead, detail) } @@ -156,8 +159,9 @@ NULL "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\"). ", + "shard_bytes = 1.9e9), or install the development ", + "version of safetensors from the mlverse/safetensors ", + "repository on GitHub, where the fix is merged. ", "Underlying error: %s"), basename(file_path), size_bytes / 1e9, underlying) } diff --git a/R/staging_ltx23.R b/R/staging.R similarity index 71% rename from R/staging_ltx23.R rename to R/staging.R index 5c95888..b5b5cea 100644 --- a/R/staging_ltx23.R +++ b/R/staging.R @@ -1,9 +1,9 @@ #' Pinned Staging for Phase-Sequential Components #' #' Phase offloading moves each large component (transformer, -#' connectors, VAEs, vocoder) between CPU and GPU every render. From -#' pageable memory those copies run through the driver's bounce -#' buffer at a fraction of PCIe speed; page-locked (pinned) host +#' connectors, VAEs, vocoder, text encoders) between CPU and GPU every +#' render. From pageable memory those copies run through the driver's +#' bounce buffer at a fraction of PCIe speed; page-locked (pinned) host #' memory transfers by DMA at full rate. Each component's parameters #' and buffers are pinned once at load; onload swaps every tensor to #' a non-blocking GPU copy of its pinned source, and offload simply @@ -18,11 +18,15 @@ #' identical, the delta is pure transfer), so pinning breaks even on #' the second render and costs a single-render session ~2s net. On by #' default; page-locking failure falls back silently per component, -#' and \code{options(diffuseR.pin_staging = FALSE)} before -#' \code{ltx23_load_pipeline} opts out (e.g. under host memory -#' pressure, where unswappable pages turn thrashing into OOM). +#' and \code{options(diffuseR.pin_staging = FALSE)} before the loader +#' opts out (e.g. under host memory pressure, where unswappable pages +#' turn thrashing into OOM). The LTX pipeline, the Gemma3 encoder, and +#' the FLUX-family image loaders (flux1, flux2, zimage) all stage +#' pinned weights; \code{\link{recommend}} computes the RAM-aware +#' \code{pin} default per model. #' -#' @name staging_ltx23 +#' @name staging +#' @aliases staging_ltx23 NULL # Allocate pinned host memory without Tensor$pin_memory(): this torch @@ -32,7 +36,7 @@ NULL # load). torch_empty_strided is the one creation op that exposes # pin_memory, so pin by allocating and copying; fall back to the noisy # path on builds where that fails. -.ltx23_pin_host <- function(p) { +.pin_host <- function(p) { tryCatch({ sz <- as.integer(p$shape) st <- if (length(sz)) { @@ -52,19 +56,24 @@ NULL #' Pin a component's tensors for fast phase transfer #' #' @param module An nn_module on the CPU. +#' @param extra Optional list of additional plain-field tensors to pin +#' alongside the module's parameters and buffers (e.g. an fp8 +#' linear's \code{weight_fp8}/\code{weight_scale} fields, which live +#' outside \code{parameters}/\code{buffers}). \code{set_data} mutates +#' each tensor in place, so the field reference stays valid. #' #' @return A list of \code{list(live, pinned)} tensor pairs, or NULL #' if pinning is unavailable (no CUDA, or page-locking failed). #' #' @keywords internal -.ltx23_pin_component <- function(module) { +.pin_component <- function(module, extra = NULL) { if (!torch::cuda_is_available()) { return(NULL) } tryCatch({ - tensors <- c(module$parameters, module$buffers) + tensors <- c(module$parameters, module$buffers, extra) lapply(tensors, function(p) { - pinned <- .ltx23_pin_host(p) + pinned <- .pin_host(p) p$set_data(pinned) list(live = p, pinned = pinned) }) @@ -77,7 +86,7 @@ NULL #' so later kernels are ordered after them; no explicit sync needed. #' #' @keywords internal -.ltx23_staged_onload <- function(staging, device) { +.staged_onload <- function(staging, device) { for (pair in staging) { pair$live$set_data(pair$pinned$to(device = device, non_blocking = TRUE)) } @@ -90,7 +99,7 @@ NULL #' are still current: offload is a pointer swap, no transfer. #' #' @keywords internal -.ltx23_staged_offload <- function(staging) { +.staged_offload <- function(staging) { for (pair in staging) { pair$live$set_data(pair$pinned) } diff --git a/R/tokenizer_bpe.R b/R/tokenizer_bpe.R index edaa862..ce31a45 100644 --- a/R/tokenizer_bpe.R +++ b/R/tokenizer_bpe.R @@ -106,6 +106,10 @@ bpe_tokenizer <- function(tokenizer_path) { #' Print BPE Tokenizer #' @param x A bpe_tokenizer object. #' @param ... Additional arguments (ignored). +#' +#' @return Invisibly returns \code{x}. Called for the side effect of +#' printing a summary of the tokenizer to the console. +#' #' @export print.bpe_tokenizer <- function(x, ...) { cat("BPE Tokenizer\n") diff --git a/R/txt2img_flux.R b/R/txt2img_flux.R index 33855ac..17416ad 100644 --- a/R/txt2img_flux.R +++ b/R/txt2img_flux.R @@ -84,19 +84,26 @@ flux_unpack_latents <- function(latents, height, width, vae_scale_factor = 8L) { #' @param device Character. Compute device. #' @param precision "nf4" or "fp8"; NULL picks the #' \code{\link{flux_memory_profile}} recommendation. -#' @param text_device Device for the text encoders ("cpu" default; the -#' T5-XXL runs float32 there). +#' @param text_device Where the text encoders compute. NULL (default) +#' takes the \code{\link{flux_memory_profile}} recommendation: "cuda" +#' on 14 GB+ cards (T5-XXL onloads in bfloat16 for its encode) and +#' "cpu" below that (T5-XXL runs float32 in place, since its ~9.8 GB +#' GPU encode phase would not fit). #' @param attn_chunk Integer or NULL. Attention query-chunk override. #' @param phase_offload Logical. One GPU tenant per phase. +#' @param pin Logical or NULL. Page-lock the phase-swapped weights for +#' DMA-rate transfer (see \code{\link{staging}}). NULL (default) +#' resolves via \code{options(diffuseR.pin_staging)} then the +#' host-RAM-aware \code{\link{recommend}} decision. #' @param verbose Logical. #' #' @return A \code{flux_pipeline} list. #' #' @export flux_load_pipeline <- function(model_dir = NULL, device = "cuda", - precision = NULL, text_device = "cpu", + precision = NULL, text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, - verbose = TRUE) { + pin = NULL, verbose = TRUE) { profile <- flux_memory_profile() if (verbose && !is.null(profile$note)) { message(profile$note) @@ -120,6 +127,16 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", } else { flux_open_checkpoint(model_dir) } + pin <- .resolve_pin(pin, "flux1", ckpt$format %||% "full") + if (is.null(text_device)) { + # The recommend tier GPU-encodes T5 only where its ~9.8 GB phase + # fits (fp8/bf16 tiers); nf4 tiers keep the CPU-fp32 encode. + text_device <- if (device == "cuda") { + profile$text_device %||% "cpu" + } else { + "cpu" + } + } if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { # Must be set before the first CUDA allocation (see @@ -137,6 +154,11 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", } else { footprint <- 4 } + if (identical(text_device, "cuda")) { + # GPU T5 encode is the largest phase (~9.8 GB); size the gc + # gate to it, not the smaller transformer footprint + footprint <- max(footprint, 10) + } .flux_gc_gates(footprint_gb = footprint) } @@ -160,7 +182,7 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", } pipe$transformer <- flux_load_transformer(ckpt, device = component_device, dtype = if (device == "cpu") "float32" else "bfloat16", - pin = device == "cuda", verbose = verbose) + pin = pin && device == "cuda", verbose = verbose) vae_config <- jsonlite::fromJSON(.flux1_cached("vae/config.json")) pipe$vae_scaling_factor <- vae_config$scaling_factor %||% 0.3611 @@ -176,9 +198,23 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", verbose = verbose ) dec$to(device = component_device) + if (device == "cuda" && phase_offload) { + # Cast to compute dtype once so the per-generation onload is + # device-only and the pinned copy is bf16 (see flux2 loader). + dec$to(dtype = torch::torch_bfloat16()) + } dec$eval() pipe$decoder <- dec + # Text encoders load to the host when phase-offloading (so they can + # stage and swap in for the encode); text_device controls where they + # compute. On the CPU-encode tiers they stay resident and compute in + # place (CLIP fp32, T5 fp32); when they GPU-encode, T5 loads bf16. + if (phase_offload) { + te_device <- "cpu" + } else { + te_device <- text_device + } if (verbose) { message("Loading CLIP text encoder...") } @@ -187,7 +223,7 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", clip, .flux1_cached("text_encoder/model.safetensors"), verbose = verbose ) - clip$to(device = text_device) + clip$to(device = te_device) clip$eval() pipe$text_encoder <- clip @@ -196,7 +232,7 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", } t5_dir <- dirname(.flux1_cached("text_encoder_2/config.json")) pipe$text_encoder2 <- load_t5_text_encoder( - t5_dir, device = text_device, + t5_dir, device = te_device, dtype = if (text_device == "cpu") "float32" else "bfloat16", verbose = verbose ) @@ -208,6 +244,13 @@ flux_load_pipeline <- function(model_dir = NULL, device = "cuda", ) pipe$scheduler_shift <- sched_cfg$shift %||% 1.0 + components <- c("transformer", "decoder") + if (!identical(text_device, "cpu")) { + components <- c(components, "text_encoder", "text_encoder2") + } + pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, + components, verbose = verbose) + structure(pipe, class = "flux_pipeline") } @@ -313,15 +356,31 @@ txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L, } phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" - onload <- function(module) { + staging <- pipeline$staging %||% list() + # Components move by name so pinned staging (see staging.R) can carry + # the CPU<->GPU transfer when the loader prepared it; otherwise the + # pageable module$to() path runs. + onload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = device) + st <- staging[[what]] + if (is.null(st)) { + module$to(device = device) + } else { + .staged_onload(st, device) + } } module } - offload <- function(module) { + offload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = "cpu") + st <- staging[[what]] + if (is.null(st)) { + module$to(device = "cpu") + } else { + .staged_offload(st) + } clear_vram() } invisible(module) @@ -330,20 +389,37 @@ txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L, t0 <- Sys.time() # --- Phase 1: text encoding ------------------------------------------------ + # On the CPU-encode tiers (text_device == "cpu") the encoders compute + # in place; otherwise each swaps to the GPU for its own encode. T5 and + # CLIP are independent (either embed may be supplied precomputed), so + # each onloads inside its own guard. + gpu_encode <- !identical(pipeline$text_device, "cpu") torch::with_no_grad({ if (is.null(prompt_embeds)) { if (verbose) { message("Encoding prompt (T5)...") } + if (gpu_encode) { + onload("text_encoder2") + } prompt_embeds <- encode_with_t5(prompt, pipeline$text_encoder2, pipeline$tokenizer2, max_sequence_length = max_sequence_length) + if (gpu_encode) { + offload("text_encoder2") + } } if (is.null(pooled_prompt_embeds)) { + if (gpu_encode) { + onload("text_encoder") + } tokens <- CLIPTokenizer(prompt) clip_device <- pipeline$text_encoder$token_embedding$weight$device tokens <- tokens$to(device = clip_device) hidden <- pipeline$text_encoder(tokens) pooled_prompt_embeds <- clip_pooled_output(hidden, tokens) + if (gpu_encode) { + offload("text_encoder") + } } }) prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) @@ -380,7 +456,7 @@ txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L, ) # --- Phase 3: denoise -------------------------------------------------------- - transformer <- onload(pipeline$transformer) + transformer <- onload("transformer") if (verbose) { message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width, height)) @@ -390,7 +466,7 @@ txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L, pooled_prompt_embeds, rope, compute_dtype, chunk_size = pipeline$attn_chunk, verbose = level ) - offload(pipeline$transformer) + offload("transformer") ltx23_release_dequant_buffers() # --- Phase 4: decode ----------------------------------------------------------- @@ -401,17 +477,14 @@ txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L, latents <- latents$div(pipeline$vae_scaling_factor %||% 0.3611)$ add(pipeline$vae_shift_factor %||% 0.1159) - decoder <- pipeline$decoder - if (phase_offload) { - decoder$to(device = device, dtype = compute_dtype) - } + decoder <- onload("decoder") torch::with_no_grad({ dec_param <- decoder$conv_in$weight img <- decoder(latents$to(device = dec_param$device, dtype = dec_param$dtype)) img <- img$to(dtype = f32)$cpu() }) - offload(decoder) + offload("decoder") img <- img$squeeze(1)$permute(c(2L, 3L, 1L)) img <- img$add(1)$div(2)$clamp(0, 1) diff --git a/R/txt2img_flux2.R b/R/txt2img_flux2.R index 25a824a..a0da374 100644 --- a/R/txt2img_flux2.R +++ b/R/txt2img_flux2.R @@ -42,6 +42,10 @@ NULL #' \code{device}; it encodes in its own phase and offloads). #' @param attn_chunk Integer or NULL. Attention query-chunk override. #' @param phase_offload Logical. One GPU tenant per phase. +#' @param pin Logical or NULL. Page-lock the phase-swapped weights for +#' DMA-rate transfer (see \code{\link{staging}}). NULL (default) +#' resolves via \code{options(diffuseR.pin_staging)} then the +#' host-RAM-aware \code{\link{recommend}} decision. #' @param verbose Logical. #' #' @return A \code{flux2_pipeline} list. @@ -50,7 +54,8 @@ NULL flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4"), text_device = NULL, attn_chunk = NULL, - phase_offload = TRUE, verbose = TRUE) { + phase_offload = TRUE, pin = NULL, + verbose = TRUE) { precision <- .flux_resolve_precision(match.arg(precision), file.path(tools::R_user_dir("diffuseR", "data"), "flux2-klein-4b-")) if (is.null(text_device)) { @@ -70,6 +75,7 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", } else { flux_open_checkpoint(model_dir) } + pin <- .resolve_pin(pin, "flux2", ckpt$format %||% "full") if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { # Resident fp8 has an NF4-like stable footprint: the native @@ -126,6 +132,12 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", verbose = verbose ) pipe$decoder$to(device = component_device) + if (device == "cuda" && phase_offload) { + # Cast to the compute dtype once here so the per-generation + # onload is device-only (and the pinned copy is bf16, not fp32); + # the decode reads the decoder's own dtype and follows. + pipe$decoder$to(dtype = torch::torch_bfloat16()) + } if (verbose) { message("Loading Qwen3 text encoder...") @@ -138,6 +150,13 @@ flux2_load_pipeline <- function(model_dir = NULL, device = "cuda", ) pipe$tokenizer <- qwen_bpe_tokenizer(.flux2_cached("tokenizer/tokenizer.json")) + components <- c("transformer", "decoder") + if (!identical(text_device, "cpu")) { + components <- c(components, "text_encoder") + } + pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, + components, verbose = verbose) + structure(pipe, class = "flux2_pipeline") } @@ -237,15 +256,31 @@ txt2img_flux2 <- function(prompt, pipeline = NULL, width = 1024L, } phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" - onload <- function(module) { + staging <- pipeline$staging %||% list() + # Components move by name so pinned staging (see staging.R) can carry + # the CPU<->GPU transfer when the loader prepared it; otherwise the + # pageable module$to() path runs. + onload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = device) + st <- staging[[what]] + if (is.null(st)) { + module$to(device = device) + } else { + .staged_onload(st, device) + } } module } - offload <- function(module) { + offload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = "cpu") + st <- staging[[what]] + if (is.null(st)) { + module$to(device = "cpu") + } else { + .staged_offload(st) + } clear_vram() } invisible(module) @@ -258,12 +293,19 @@ txt2img_flux2 <- function(prompt, pipeline = NULL, width = 1024L, if (verbose) { message("Encoding prompt (Qwen3)...") } - onload(pipeline$text_encoder) + # An explicit text_device = "cpu" keeps Qwen3 fp32 on the host + # (it computes in place); only onload it when it phase-swaps. + gpu_encode <- !identical(pipeline$text_device, "cpu") + if (gpu_encode) { + onload("text_encoder") + } te_device <- pipeline$text_encoder$model$embed_tokens$weight$device prompt_embeds <- encode_with_qwen3(prompt, pipeline$text_encoder, pipeline$tokenizer, max_sequence_length = max_sequence_length, device = te_device) - offload(pipeline$text_encoder) + if (gpu_encode) { + offload("text_encoder") + } } prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) txt_len <- prompt_embeds$shape[2] @@ -301,8 +343,13 @@ txt2img_flux2 <- function(prompt, pipeline = NULL, width = 1024L, ) # --- Phase 3: denoise ------------------------------------------------------------ - transformer <- onload(pipeline$transformer) - if (isTRUE(pipeline$fp8_resident)) { + transformer <- onload("transformer") + # Resident fp8's plain weight fields ride to the GPU with the phase. + # When staging covers the transformer it already moved them (and the + # field reassignment here would orphan the staged pairs), so skip it. + fp8_manual <- isTRUE(pipeline$fp8_resident) && + is.null(staging[["transformer"]]) + if (fp8_manual) { .flux_fp8_to_device(transformer, device) } if (verbose) { @@ -314,10 +361,10 @@ txt2img_flux2 <- function(prompt, pipeline = NULL, width = 1024L, compute_dtype, chunk_size = pipeline$attn_chunk, verbose = level ) - if (isTRUE(pipeline$fp8_resident) && phase_offload) { + if (fp8_manual && phase_offload) { .flux_fp8_to_device(pipeline$transformer, "cpu") } - offload(pipeline$transformer) + offload("transformer") ltx23_release_dequant_buffers() # --- Phase 4: decode --------------------------------------------------------------- @@ -332,17 +379,14 @@ txt2img_flux2 <- function(prompt, pipeline = NULL, width = 1024L, ) latents <- flux2_unpatchify_latents(latents) - decoder <- pipeline$decoder - if (phase_offload) { - decoder$to(device = device, dtype = compute_dtype) - } + decoder <- onload("decoder") torch::with_no_grad({ dec_param <- decoder$post_quant_conv$weight img <- decoder(latents$to(device = dec_param$device, dtype = dec_param$dtype)) img <- img$to(dtype = f32)$cpu() }) - offload(decoder) + offload("decoder") img <- img$squeeze(1)$permute(c(2L, 3L, 1L)) img <- img$add(1)$div(2)$clamp(0, 1) diff --git a/R/txt2img_zimage.R b/R/txt2img_zimage.R index 859113f..482e385 100644 --- a/R/txt2img_zimage.R +++ b/R/txt2img_zimage.R @@ -43,6 +43,10 @@ NULL #' \code{device}; it encodes in its own phase and offloads). #' @param attn_chunk Integer or NULL. Attention query-chunk override. #' @param phase_offload Logical. One GPU tenant per phase. +#' @param pin Logical or NULL. Page-lock the phase-swapped weights for +#' DMA-rate transfer (see \code{\link{staging}}). NULL (default) +#' resolves via \code{options(diffuseR.pin_staging)} then the +#' host-RAM-aware \code{\link{recommend}} decision. #' @param verbose Logical. #' #' @return A \code{zimage_pipeline} list. @@ -51,7 +55,8 @@ NULL zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4"), text_device = NULL, attn_chunk = NULL, - phase_offload = TRUE, verbose = TRUE) { + phase_offload = TRUE, pin = NULL, + verbose = TRUE) { precision <- .flux_resolve_precision(match.arg(precision), file.path(tools::R_user_dir("diffuseR", "data"), "zimage-turbo-")) if (is.null(text_device)) { @@ -71,6 +76,7 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", } else { flux_open_checkpoint(model_dir) } + pin <- .resolve_pin(pin, "zimage", ckpt$format %||% "full") if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) { # Stable resident footprint: the native backend avoids @@ -126,6 +132,11 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", verbose = verbose ) dec$to(device = component_device) + if (device == "cuda" && phase_offload) { + # Cast to compute dtype once so the per-generation onload is + # device-only and the pinned copy is bf16 (see flux2 loader). + dec$to(dtype = torch::torch_bfloat16()) + } dec$eval() pipe$decoder <- dec @@ -146,6 +157,13 @@ zimage_load_pipeline <- function(model_dir = NULL, device = "cuda", ) pipe$tokenizer <- qwen_bpe_tokenizer(.zimage_cached("tokenizer/tokenizer.json")) + components <- c("transformer", "decoder") + if (!identical(text_device, "cpu")) { + components <- c(components, "text_encoder") + } + pipe$staging <- .flux_build_staging(pipe, pin, phase_offload, device, + components, verbose = verbose) + structure(pipe, class = "zimage_pipeline") } @@ -262,15 +280,31 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, } phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu" - onload <- function(module) { + staging <- pipeline$staging %||% list() + # Components move by name so pinned staging (see staging.R) can carry + # the CPU<->GPU transfer when the loader prepared it; otherwise the + # pageable module$to() path runs. + onload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = device) + st <- staging[[what]] + if (is.null(st)) { + module$to(device = device) + } else { + .staged_onload(st, device) + } } module } - offload <- function(module) { + offload <- function(what) { + module <- pipeline[[what]] if (phase_offload) { - module$to(device = "cpu") + st <- staging[[what]] + if (is.null(st)) { + module$to(device = "cpu") + } else { + .staged_offload(st) + } clear_vram() } invisible(module) @@ -283,13 +317,20 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, if (verbose) { message("Encoding prompt (Qwen3)...") } - onload(pipeline$text_encoder) + # An explicit text_device = "cpu" keeps Qwen3 fp32 on the host + # (it computes in place); only onload it when it phase-swaps. + gpu_encode <- !identical(pipeline$text_device, "cpu") + if (gpu_encode) { + onload("text_encoder") + } te_device <- pipeline$text_encoder$model$embed_tokens$weight$device prompt_embeds <- .zimage_encode_prompt(prompt, pipeline$text_encoder, pipeline$tokenizer, penult_layer = pipeline$te_penult_layer %||% 35L, max_sequence_length = max_sequence_length, device = te_device) - offload(pipeline$text_encoder) + if (gpu_encode) { + offload("text_encoder") + } } prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype) @@ -313,8 +354,13 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, ) # --- Phase 3: denoise ------------------------------------------------------------ - transformer <- onload(pipeline$transformer) - if (isTRUE(pipeline$fp8_resident)) { + transformer <- onload("transformer") + # Resident fp8's plain weight fields ride to the GPU with the phase. + # When staging covers the transformer it already moved them (and the + # field reassignment here would orphan the staged pairs), so skip it. + fp8_manual <- isTRUE(pipeline$fp8_resident) && + is.null(staging[["transformer"]]) + if (fp8_manual) { .flux_fp8_to_device(transformer, device) } if (verbose) { @@ -326,10 +372,10 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, compute_dtype, chunk_size = pipeline$attn_chunk, verbose = level ) - if (isTRUE(pipeline$fp8_resident) && phase_offload) { + if (fp8_manual && phase_offload) { .flux_fp8_to_device(pipeline$transformer, "cpu") } - offload(pipeline$transformer) + offload("transformer") ltx23_release_dequant_buffers() # --- Phase 4: decode --------------------------------------------------------------- @@ -339,17 +385,14 @@ txt2img_zimage <- function(prompt, pipeline = NULL, width = 1024L, latents <- latents$div(pipeline$vae_scaling_factor %||% 0.3611)$ add(pipeline$vae_shift_factor %||% 0.1159) - decoder <- pipeline$decoder - if (phase_offload) { - decoder$to(device = device, dtype = compute_dtype) - } + decoder <- onload("decoder") torch::with_no_grad({ dec_param <- decoder$conv_in$weight img <- decoder(latents$to(device = dec_param$device, dtype = dec_param$dtype)) img <- img$to(dtype = f32)$cpu() }) - offload(decoder) + offload("decoder") img <- img$squeeze(1)$permute(c(2L, 3L, 1L)) img <- img$add(1)$div(2)$clamp(0, 1) diff --git a/R/txt2vid_ltx23.R b/R/txt2vid_ltx23.R index 698ea0e..8501d7d 100644 --- a/R/txt2vid_ltx23.R +++ b/R/txt2vid_ltx23.R @@ -352,7 +352,7 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", isTRUE(getOption("diffuseR.pin_staging", TRUE))) { # Page-lock every phase-offloaded component once so the # per-render CPU<->GPU moves run at full PCIe rate (offload - # becomes a pointer swap; see staging_ltx23.R). Falls back + # becomes a pointer swap; see staging.R). Falls back # silently per component if page-locking fails. if (verbose) { message("Pinning host staging buffers...") @@ -360,7 +360,7 @@ ltx23_load_pipeline <- function(checkpoint_path, device = "cuda", staging <- list() for (nm in intersect(c("transformer", "connectors", "vae", "audio_vae", "vocoder"), names(pipe))) { - st <- .ltx23_pin_component(pipe[[nm]]) + st <- .pin_component(pipe[[nm]]) if (!is.null(st)) { staging[[nm]] <- st } @@ -568,7 +568,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, # Each phase is the sole GPU tenant: components move on for their # phase and back off afterwards. Pipeline components are referred - # to by name so pinned staging (see staging_ltx23.R) can be used + # to by name so pinned staging (see staging.R) can be used # when the loader prepared it; plain modules (the upsampler) take # the pageable path. phase_offload <- phase_offload && device != "cpu" @@ -612,7 +612,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, if (is.null(st)) { module$to(device = device) } else { - .ltx23_staged_onload(st, device) + .staged_onload(st, device) } } module @@ -635,7 +635,7 @@ txt2vid_ltx2 <- function(prompt, pipeline, text_encoder = NULL, if (is.null(st)) { module$to(device = "cpu") } else { - .ltx23_staged_offload(st) + .staged_offload(st) } # gc only -- NO cuda_empty_cache between phases: returning # blocks to the driver forces the next phase to regrow the diff --git a/R/upsampler_ltx23.R b/R/upsampler_ltx23.R index 9fff159..bfc7f16 100644 --- a/R/upsampler_ltx23.R +++ b/R/upsampler_ltx23.R @@ -36,6 +36,10 @@ ltx23_upsampler_res_block <- torch::nn_module( #' @param mid_channels Integer. #' @param num_blocks_per_stage Integer. #' +#' @return Module whose forward(hidden_states) returns the 2x +#' spatially upscaled latent, a tensor with the same batch, channel +#' and frame counts and doubled height and width. +#' #' @export ltx23_latent_upsampler <- torch::nn_module( "ltx23_latent_upsampler", diff --git a/R/vae_decoder.R b/R/vae_decoder.R index d41ba90..17731c8 100644 --- a/R/vae_decoder.R +++ b/R/vae_decoder.R @@ -315,11 +315,21 @@ load_decoder_weights <- function(native_decoder, torchscript_path, #' @export #' #' @examples +#' if (torch::torch_is_installed()) { +#' # A small decoder; the SD/SDXL defaults are far too large to build +#' # inside an example. +#' decoder <- vae_decoder_native(latent_channels = 4, +#' block_channels = 32, +#' norm_groups = 32) +#' latents <- torch::torch_randn(c(1, 4, 8, 8)) +#' image <- torch::with_no_grad(decoder(latents)) +#' image$shape +#' } +#' +#' # Real weights come from a downloaded checkpoint. #' \dontrun{ #' decoder <- vae_decoder_native() #' load_decoder_weights(decoder, "path/to/decoder.pt") -#' latents <- torch::torch_randn(c(1, 4, 64, 64)) -#' image <- decoder(latents) #' } vae_decoder_native <- torch::nn_module( "VAEDecoderNative", diff --git a/R/vae_ltx23.R b/R/vae_ltx23.R index f9d2201..83b7704 100644 --- a/R/vae_ltx23.R +++ b/R/vae_ltx23.R @@ -26,6 +26,10 @@ NULL #' @param is_causal Logical. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' encoded video latent [B, C, F, H, W], with the last channel +#' repeated to carry the per-channel scale expected downstream. +#' #' @export ltx23_video_encoder3d <- torch::nn_module( "ltx23_video_encoder3d", @@ -139,6 +143,10 @@ ltx23_video_encoder3d <- torch::nn_module( #' @param upsample_factor Integer vector per up block. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' decoded pixel tensor [B, 3, F, H, W], with the final patch axes +#' flattened back into height and width. +#' #' @export ltx23_video_decoder3d <- torch::nn_module( "ltx23_video_decoder3d", @@ -256,6 +264,10 @@ ltx23_video_decoder3d <- torch::nn_module( #' @param encoder_causal,decoder_causal Logicals. Temporal padding modes. #' @param encoder_spatial_padding_mode,decoder_spatial_padding_mode Characters. #' +#' @return Module bundling the video encoder and decoder. Its +#' forward(z) is \code{decode(z)}, returning pixels for a latent; +#' \code{$encode()} and \code{$decode()} are callable separately. +#' #' @export ltx23_video_vae <- torch::nn_module( "ltx23_video_vae", diff --git a/R/vae_ltx23_modules.R b/R/vae_ltx23_modules.R index c23acce..6782b50 100644 --- a/R/vae_ltx23_modules.R +++ b/R/vae_ltx23_modules.R @@ -17,6 +17,9 @@ NULL #' #' @param eps Numeric. Stability epsilon. #' +#' @return Module whose forward(x) returns \code{x} divided by its +#' per-channel root mean square, a tensor of the same shape. +#' #' @export ltx23_per_channel_rms_norm <- torch::nn_module( "ltx23_per_channel_rms_norm", @@ -40,6 +43,11 @@ ltx23_per_channel_rms_norm <- torch::nn_module( #' @param stride Integer or length-3 vector. #' @param spatial_padding_mode Character. Conv padding mode. #' +#' @return Module whose forward(hidden_states, causal) returns the 3-D +#' convolution of the input. With \code{causal = TRUE} the temporal +#' axis is left-padded by replicating the first frame, so no output +#' frame sees a later input frame. +#' #' @export ltx23_causal_conv3d <- torch::nn_module( "ltx23_causal_conv3d", @@ -91,6 +99,9 @@ ltx23_causal_conv3d <- torch::nn_module( #' @param eps Numeric. Shortcut LayerNorm epsilon. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(inputs, causal) returns \code{inputs} +#' plus the residual branch, a tensor of the same shape. +#' #' @export ltx23_video_resnet_block3d <- torch::nn_module( "ltx23_video_resnet_block3d", @@ -150,6 +161,10 @@ ltx23_video_resnet_block3d <- torch::nn_module( #' @param stride Length-3 integer vector (t, h, w). #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' space-to-depth downsampled tensor plus its residual: spatial and +#' temporal extents shrink by \code{stride}, channels grow to match. +#' #' @export ltx23_video_downsampler3d <- torch::nn_module( "ltx23_video_downsampler3d", @@ -206,6 +221,10 @@ ltx23_video_downsampler3d <- torch::nn_module( #' @param upscale_factor Integer. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' depth-to-space upsampled tensor: spatial and temporal extents grow +#' by \code{stride}, channels shrink to match. +#' #' @export ltx23_video_upsampler3d <- torch::nn_module( "ltx23_video_upsampler3d", @@ -272,6 +291,10 @@ ltx23_video_upsampler3d <- torch::nn_module( #' @param downsample_type "spatial", "temporal", or "spatiotemporal". #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' stage output: the resnet stack applied in sequence, then the +#' optional downsampler. +#' #' @export ltx23_video_down_block3d <- torch::nn_module( "ltx23_video_down_block3d", @@ -329,6 +352,9 @@ ltx23_video_down_block3d <- torch::nn_module( #' @param resnet_eps Numeric. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' bottleneck output, a tensor of the same shape as the input. +#' #' @export ltx23_video_mid_block3d <- torch::nn_module( "ltx23_video_mid_block3d", @@ -366,6 +392,10 @@ ltx23_video_mid_block3d <- torch::nn_module( #' @param upscale_factor Integer. #' @param spatial_padding_mode Character. #' +#' @return Module whose forward(hidden_states, causal) returns the +#' stage output: the optional input projection and upsampler, then the +#' resnet stack applied in sequence. +#' #' @export ltx23_video_up_block3d <- torch::nn_module( "ltx23_video_up_block3d", diff --git a/R/vocoder_ltx23.R b/R/vocoder_ltx23.R index 765b2da..2a39c4b 100644 --- a/R/vocoder_ltx23.R +++ b/R/vocoder_ltx23.R @@ -74,6 +74,9 @@ ltx23_kaiser_sinc_filter1d <- function(cutoff, half_width, kernel_size) { #' @param ratio Integer. Downsampling ratio. #' @param kernel_size Integer or NULL (default 6*ratio rounded even). #' +#' @return Module whose forward(x) returns \code{x} low-pass filtered +#' and decimated by \code{ratio} along the time axis. +#' #' @export ltx23_downsample1d <- torch::nn_module( "ltx23_downsample1d", @@ -102,6 +105,10 @@ ltx23_downsample1d <- torch::nn_module( #' @param persistent Logical. Register the filter as a buffer (present in #' checkpoints); FALSE stores the computed filter as a plain field. #' +#' @return Module whose forward(x) returns \code{x} interpolated up by +#' \code{ratio} along the time axis, with the filter padding trimmed +#' off. +#' #' @export ltx23_upsample1d <- torch::nn_module( "ltx23_upsample1d", @@ -164,6 +171,10 @@ ltx23_upsample1d <- torch::nn_module( #' @param channels Integer. #' @param eps Numeric. #' +#' @return Module whose forward(hidden_states) returns the Snake +#' activation \code{x + sin(alpha * x)^2 / beta}, a tensor of the same +#' shape as the input. +#' #' @export ltx23_snake_beta <- torch::nn_module( "ltx23_snake_beta", @@ -189,6 +200,11 @@ ltx23_snake_beta <- torch::nn_module( #' @param channels Integer. Channels for the SnakeBeta activation. #' @param ratio,kernel_size Integers. Resampling config. #' +#' @return Module whose forward(x) returns the activation applied at 2x +#' rate (upsample, activate, downsample), a tensor of the same shape +#' as \code{x}, with the aliasing the raw activation would introduce +#' filtered out. +#' #' @export ltx23_antialias_act1d <- torch::nn_module( "ltx23_antialias_act1d", @@ -213,6 +229,10 @@ ltx23_antialias_act1d <- torch::nn_module( #' @param dilations Integer vector. #' @param antialias_ratio,antialias_kernel_size Integers. #' +#' @return Module whose forward(x) returns \code{x} after the dilated +#' convolution pairs have been added back as residuals, a tensor of +#' the same shape. +#' #' @export ltx23_vocoder_resblock <- torch::nn_module( "ltx23_vocoder_resblock", @@ -266,6 +286,9 @@ ltx23_vocoder_resblock <- torch::nn_module( #' @param antialias_ratio,antialias_kernel_size Integers. #' @param final_bias Logical. #' +#' @return Module whose forward(hidden_states, time_last) returns the +#' synthesized waveform [B, 1, samples] for a mel spectrogram. +#' #' @export ltx23_vocoder <- torch::nn_module( "ltx23_vocoder", @@ -368,6 +391,9 @@ ltx23_causal_stft <- torch::nn_module( #' #' @param filter_length,hop_length,window_length,num_mel_channels Integers. #' +#' @return Module whose forward(waveform) returns the log-mel +#' spectrogram [B, n_mels, frames], clamped at 1e-5 before the log. +#' #' @export ltx23_mel_stft <- torch::nn_module( "ltx23_mel_stft", @@ -407,6 +433,10 @@ ltx23_mel_stft <- torch::nn_module( #' @param input_sampling_rate,output_sampling_rate Integers. #' @param hop_length Integer. Mel analysis hop. #' +#' @return Module whose forward(mel_spec) returns the +#' bandwidth-extended waveform [B, 1, samples], trimmed to the sample +#' count implied by the input frames and the rate ratio. +#' #' @export ltx23_vocoder_with_bwe <- torch::nn_module( "ltx23_vocoder_with_bwe", diff --git a/R/vram.R b/R/vram.R index b254934..bdaf55a 100644 --- a/R/vram.R +++ b/R/vram.R @@ -15,11 +15,9 @@ NULL #' @export #' #' @examples -#' \dontrun{ -#' if (is_blackwell_gpu()) { -#' message("Using Blackwell-compatible settings") -#' } -#' } +#' # Soft probe: FALSE on any machine without a Blackwell card, and on +#' # machines where torch has no lantern binaries. +#' is_blackwell_gpu() is_blackwell_gpu <- function() { # Check compute capability via torch. cuda_is_available() ERRORS # (not FALSE) when the torch package is installed without its @@ -83,10 +81,9 @@ is_blackwell_gpu <- function() { #' @export #' #' @examples -#' \dontrun{ -#' model$to(device = "cuda") -#' output <- model(x) -#' offload_to_cpu(model) +#' if (torch::torch_is_installed()) { +#' model <- torch::nn_linear(4, 2) +#' offload_to_cpu(model) #' } offload_to_cpu <- function(module, gc = TRUE) { module$to(device = "cpu") @@ -109,10 +106,10 @@ offload_to_cpu <- function(module, gc = TRUE) { #' @export #' #' @examples -#' \dontrun{ -#' load_to_gpu(model) -#' output <- model(x) -#' offload_to_cpu(model) +#' if (torch::torch_is_installed()) { +#' model <- torch::nn_linear(4, 2) +#' # "cuda" needs a GPU; "cpu" is the portable round trip. +#' load_to_gpu(model, device = "cpu") #' } load_to_gpu <- function(module, device = "cuda") { module$to(device = device) @@ -130,8 +127,8 @@ load_to_gpu <- function(module, device = "cuda") { #' @export #' #' @examples -#' \dontrun{ -#' vram_report("After model load") +#' if (torch::torch_is_installed()) { +#' vram_report("After model load") #' } vram_report <- function(label = "") { if (!torch::cuda_is_available()) { @@ -170,8 +167,8 @@ vram_report <- function(label = "") { #' @export #' #' @examples -#' \dontrun{ -#' clear_vram() +#' if (torch::torch_is_installed()) { +#' clear_vram() #' } clear_vram <- function(verbose = FALSE) { if (!torch::cuda_is_available()) { diff --git a/cran-comments.md b/cran-comments.md index 421781d..8ec5605 100644 --- a/cran-comments.md +++ b/cran-comments.md @@ -20,6 +20,30 @@ package. ## Notes for the reviewers +* On the previous submission we were asked to replace `\dontrun{}` + with `\donttest{}`. We unwrapped 14 of the 23 examples so they now + run during check (device/profile policy helpers, both schedulers, + the VRAM helpers, `save_image()`, `save_video()` frame output, and a + small `vae_decoder_native()`), and rewrote them to be self-contained + rather than referencing undefined objects. Nine remain under + `\dontrun{}` because they cannot execute anywhere without model + weights on disk: + + - `download_model()`, `download_component()` - network downloads of + multi-GB weights. + - `load_pipeline()`, `load_model_component()` - require those + downloaded weights. + - `txt2img()`, `txt2img_sd21()`, `txt2img_sdxl()` - full image + generation; needs the weights and minutes of compute. + - `ltx23_open_checkpoint()` - needs a ~20 GB LTX-2.3 checkpoint the + user downloads under Lightricks' own license. + - `load_decoder_weights()` in the `vae_decoder_native()` page - + needs a checkpoint file. The rest of that example runs. + + `\donttest{}` would not help for these: CRAN runs `\donttest{}` + examples, and each of these would either attempt a large download or + fail on a missing file. + * Model weights are never bundled and never downloaded without consent: every download function prompts interactively with the size stated, and non-interactive sessions require diff --git a/inst/tinytest/test_flux_staging.R b/inst/tinytest/test_flux_staging.R new file mode 100644 index 0000000..6e7bbc0 --- /dev/null +++ b/inst/tinytest/test_flux_staging.R @@ -0,0 +1,118 @@ +# Flux-family pinned staging surface + gate logic. The CPU/no-torch +# parts (formals, .resolve_pin precedence, .flux_build_staging opt-out +# paths) run everywhere; the pin round trip is CUDA-gated below. + +library(diffuseR) + +# --- loader surface: pin + text_device knobs -------------------------------------- + +for (fn in list(flux_load_pipeline, flux2_load_pipeline, + zimage_load_pipeline)) { + expect_true("pin" %in% names(formals(fn))) + expect_true("text_device" %in% names(formals(fn))) +} +# flux1's text_device default moved from "cpu" to NULL (profile-resolved) +expect_null(eval(formals(flux_load_pipeline)$text_device)) + +# --- .resolve_pin precedence: explicit > option > decision ------------------------ + +# Explicit argument wins over everything. +old <- getOption("diffuseR.pin_staging") +on.exit(options(diffuseR.pin_staging = old), add = TRUE) +options(diffuseR.pin_staging = FALSE) +expect_true(diffuseR:::.resolve_pin(TRUE, "flux1", "nf4")) +expect_false(diffuseR:::.resolve_pin(FALSE, "flux1", "nf4")) +# NULL falls through to the option. +expect_false(diffuseR:::.resolve_pin(NULL, "flux1", "nf4")) +options(diffuseR.pin_staging = TRUE) +expect_true(diffuseR:::.resolve_pin(NULL, "flux1", "nf4")) +# Unset option falls through to the RAM-aware decision (logical result). +options(diffuseR.pin_staging = NULL) +expect_true(is.logical(diffuseR:::.resolve_pin(NULL, "flux2", "fp8"))) + +# --- .pin_decision: cpu -> FALSE, NA RAM -> TRUE, else 2x rule --------------------- + +expect_false(diffuseR:::.pin_decision("flux1", "nf4", host_ram_gb = 64, + cpu = TRUE)) +expect_true(diffuseR:::.pin_decision("flux1", "nf4", host_ram_gb = NA)) +expect_true(diffuseR:::.pin_decision("flux1", "nf4", host_ram_gb = 128)) +expect_false(diffuseR:::.pin_decision("flux1", "nf4", host_ram_gb = 8)) + +# --- .pinned_set_gb flux1 row: T5 host copy now bf16 (was fp32) -------------------- + +expect_equal(diffuseR:::.pinned_set_gb("flux1", "nf4"), 17) +expect_equal(diffuseR:::.pinned_set_gb("flux1", "fp8"), 22) +expect_equal(diffuseR:::.pinned_set_gb("flux1", "bf16"), 34) + +# --- flux tier text_device: flux1 GPU-encodes only where T5 fits ------------------ + +# fp8/bf16 tiers (>=14 GB) onload T5; nf4 tiers keep CPU-fp32 encode. +expect_equal(recommend("flux1", vram_gb = 24)$text_device, "cuda") +expect_equal(recommend("flux1", vram_gb = 16)$text_device, "cuda") # fp8 tier +expect_equal(recommend("flux1", vram_gb = 10)$text_device, "cpu") # nf4 tier +expect_equal(recommend("flux1", vram_gb = 0)$text_device, "cpu") # cpu tier +# The small flux2/zimage encoder rides the GPU on every GPU tier. +expect_equal(recommend("flux2", vram_gb = 16)$text_device, "cuda") + +# --- .flux_build_staging opt-out paths return NULL (no torch needed) -------------- + +dummy <- list(format = "nf4") +expect_null(diffuseR:::.flux_build_staging(dummy, pin = FALSE, + phase_offload = TRUE, device = "cuda", components = "transformer")) +expect_null(diffuseR:::.flux_build_staging(dummy, pin = TRUE, + phase_offload = FALSE, device = "cuda", components = "transformer")) +expect_null(diffuseR:::.flux_build_staging(dummy, pin = TRUE, + phase_offload = TRUE, device = "cpu", components = "transformer")) + +# --- fp8 field collector + pin round trip (torch / CUDA gated) -------------------- + +if (!requireNamespace("torch", quietly = TRUE) || + !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +# .flux_fp8_collect walks children and returns weight_fp8 + weight_scale +# pairs. Build two fp8 linears under a parent and check the count. +parent <- torch::nn_module( + "fp8_parent", + initialize = function() { + self$a <- ltx23_fp8_linear(4L, 4L, bias = FALSE) + self$b <- ltx23_fp8_linear(4L, 4L, bias = FALSE) + }, + forward = function(x) x +)() +w <- torch::torch_randn(4L, 4L)$to(dtype = torch::torch_float8_e4m3fn()) +s <- torch::torch_tensor(0.5) +parent$a$set_fp8_weight(w, s) +parent$b$set_fp8_weight(w, s) +collected <- diffuseR:::.flux_fp8_collect(parent) +expect_equal(length(collected), 4L) # 2 linears x (weight_fp8, weight_scale) + +if (!torch::cuda_is_available()) { + exit_file("no CUDA") +} + +# .pin_component with extra fields pins params/buffers AND the fp8 fields; +# staged onload/offload moves them all and preserves the forward output. +m <- ltx23_feed_forward(16L) +m$eval() +extra_t <- torch::torch_randn(8L, 8L)$to(dtype = torch::torch_float8_e4m3fn()) +scale0 <- torch::torch_tensor(2.0) # 0-dim scale must pin too +st <- diffuseR:::.pin_component(m, extra = list(extra_t, scale0)) +expect_false(is.null(st)) +# params + buffers + 2 extras all became pinned live/pinned pairs +expect_true(length(st) >= length(c(m$parameters, m$buffers)) + 2L) +expect_true(suppressWarnings( + st[[length(st)]]$live$is_pinned(device = torch::torch_device("cuda")) +)) + +x <- torch::torch_randn(2L, 5L, 16L) +torch::with_no_grad(ref <- m(x)) +diffuseR:::.staged_onload(st, "cuda") +expect_equal(st[[1]]$live$device$type, "cuda") +torch::with_no_grad(out_gpu <- m(x$to(device = "cuda"))$cpu()) +expect_true(as.numeric((out_gpu - ref)$abs()$max()) < 1e-5) +diffuseR:::.staged_offload(st) +expect_equal(st[[1]]$live$device$type, "cpu") +torch::with_no_grad(out_back <- m(x)) +expect_true(as.numeric((out_back - ref)$abs()$max()) == 0) diff --git a/inst/tinytest/test_staging_ltx23.R b/inst/tinytest/test_staging.R similarity index 82% rename from inst/tinytest/test_staging_ltx23.R rename to inst/tinytest/test_staging.R index 7345855..0490091 100644 --- a/inst/tinytest/test_staging_ltx23.R +++ b/inst/tinytest/test_staging.R @@ -1,4 +1,4 @@ -# Pinned staging round trip (R/staging_ltx23.R): pin -> onload -> +# Pinned staging round trip (R/staging.R): pin -> onload -> # offload -> onload must preserve outputs exactly. CUDA-only. if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { @@ -14,7 +14,7 @@ m$eval() x <- torch::torch_randn(2L, 5L, 16L) torch::with_no_grad(ref <- m(x)) -st <- diffuseR:::.ltx23_pin_component(m) +st <- diffuseR:::.pin_component(m) expect_false(is.null(st)) expect_true(suppressWarnings( st[[1]]$live$is_pinned(device = torch::torch_device("cuda")) @@ -25,7 +25,7 @@ torch::with_no_grad(out_pinned <- m(x)) expect_true(as.numeric((out_pinned - ref)$abs()$max()) == 0) # Onload: GPU forward matches -diffuseR:::.ltx23_staged_onload(st, "cuda") +diffuseR:::.staged_onload(st, "cuda") expect_equal(st[[1]]$live$device$type, "cuda") torch::with_no_grad( out_gpu <- m(x$to(device = "cuda"))$cpu() @@ -33,13 +33,13 @@ torch::with_no_grad( expect_true(as.numeric((out_gpu - ref)$abs()$max()) < 1e-5) # Offload: pointer swap back to the pinned copies, exact outputs -diffuseR:::.ltx23_staged_offload(st) +diffuseR:::.staged_offload(st) expect_equal(st[[1]]$live$device$type, "cpu") torch::with_no_grad(out_back <- m(x)) expect_true(as.numeric((out_back - ref)$abs()$max()) == 0) # Second round trip still exact -diffuseR:::.ltx23_staged_onload(st, "cuda") +diffuseR:::.staged_onload(st, "cuda") torch::with_no_grad(out_gpu2 <- m(x$to(device = "cuda"))$cpu()) expect_true(as.numeric((out_gpu2 - out_gpu)$abs()$max()) == 0) -diffuseR:::.ltx23_staged_offload(st) +diffuseR:::.staged_offload(st) diff --git a/inst/tinytest/test_txt2img_zimage.R b/inst/tinytest/test_txt2img_zimage.R new file mode 100644 index 0000000..d30d893 --- /dev/null +++ b/inst/tinytest/test_txt2img_zimage.R @@ -0,0 +1,62 @@ +# End-to-end smoke test for the Z-Image Turbo pipeline wiring on the CPU +# with tiny random-init components: latents -> denoise (2 steps, static +# shift, reversed timestep + negated output) -> scale/shift -> decode. +# Numeric quality comes from the per-component parity tests; this checks +# the phase plumbing (onload/offload routing, decode). + +if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) { + exit_file("torch not fully installed") +} + +library(diffuseR) + +torch::torch_manual_seed(9) + +transformer <- zimage_transformer( + in_channels = 16L, dim = 48L, n_layers = 1L, n_refiner_layers = 1L, + n_heads = 2L, cap_feat_dim = 24L, axes_dims = c(8L, 8L, 8L) +) +transformer$eval() + +decoder <- vae_decoder_native( + latent_channels = 16L, + block_channels = c(32L, 32L, 16L, 8L), + norm_groups = 8L +) +decoder$eval() + +pipeline <- structure( + list( + transformer = transformer, + decoder = decoder, + device = "cpu", + text_device = "cpu", + phase_offload = FALSE, + fp8_resident = FALSE, + format = "full", + attn_chunk = NULL, + config = list(in_channels = 16L), + sched_shift = 3.0, + te_penult_layer = 35L, + vae_scaling_factor = 0.3611, + vae_shift_factor = 0.1159 + ), + class = "zimage_pipeline" +) + +res <- txt2img_zimage( + "tiny smoke test", + pipeline = pipeline, + width = 64L, height = 64L, + num_inference_steps = 2L, + seed = 42L, + prompt_embeds = torch::torch_randn(7L, 24L), + save_file = FALSE, + verbose = FALSE +) + +expect_equal(dim(res$image), c(64L, 64L, 3L)) +expect_true(all(is.finite(res$image))) +expect_true(all(res$image >= 0) && all(res$image <= 1)) +expect_equal(res$metadata$steps, 2L) +expect_equal(res$metadata$model, "zimage-turbo") diff --git a/man/auto_devices.Rd b/man/auto_devices.Rd index 4f23c5b..42e50a5 100644 --- a/man/auto_devices.Rd +++ b/man/auto_devices.Rd @@ -32,14 +32,11 @@ compatibility issues, regardless of available VRAM. The native modules (`use_native_unet` and friends) do not have this restriction. } \examples{ -\dontrun{ -# Auto-detect best configuration -devices <- auto_devices("sdxl") +# Force a strategy: no GPU or nvidia-smi needed. +str(auto_devices("sdxl", strategy = "cpu_only")) -# Use with models2devices -m2d <- models2devices("sdxl", devices = auto_devices("sdxl")) +str(auto_devices("sd21", strategy = "unet_gpu")) -# Force CPU-only -devices <- auto_devices("sdxl", strategy = "cpu_only") -} +# Auto-detect free VRAM and pick a strategy for this machine. +str(auto_devices("sdxl")) } diff --git a/man/clear_vram.Rd b/man/clear_vram.Rd index cc2b746..858b4d5 100644 --- a/man/clear_vram.Rd +++ b/man/clear_vram.Rd @@ -15,7 +15,7 @@ Invisibly returns NULL. Forces garbage collection and clears CUDA memory cache. } \examples{ -\dontrun{ -clear_vram() +if (torch::torch_is_installed()) { + clear_vram() } } diff --git a/man/ddim_scheduler_create.Rd b/man/ddim_scheduler_create.Rd index 6cd646b..3eaf630 100644 --- a/man/ddim_scheduler_create.Rd +++ b/man/ddim_scheduler_create.Rd @@ -9,7 +9,7 @@ ddim_scheduler_create(num_train_timesteps = 1000, num_inference_steps = 50, beta_start = 0.00085, beta_end = 0.012, rescale_betas_zero_snr = FALSE, dtype = torch::torch_float32(), - device = c(torch::torch_device("cpu"), torch::torch_device("cuda"))) + device = torch::torch_device("cpu")) } \arguments{ \item{num_train_timesteps}{Integer. The number of diffusion steps used to @@ -66,13 +66,13 @@ Song, J., Meng, C., & Ermon, S. (2020). \url{https://arxiv.org/abs/2010.02502} } \examples{ -\dontrun{ -# Create a DDIM scheduler with custom parameters -scheduler <- ddim_scheduler_create( - num_train_timesteps = 1000, - num_inference_steps = 30, - eta = 0.5, - beta_schedule = "scaled_linear" -) +if (torch::torch_is_installed()) { + scheduler <- ddim_scheduler_create( + num_train_timesteps = 1000, + num_inference_steps = 5, + eta = 0.5, + beta_schedule = "scaled_linear" + ) + scheduler$timesteps } } diff --git a/man/ddim_scheduler_step.Rd b/man/ddim_scheduler_step.Rd index 9df3960..1945efb 100644 --- a/man/ddim_scheduler_step.Rd +++ b/man/ddim_scheduler_step.Rd @@ -89,13 +89,17 @@ Salimans, T., & Ho, J. (2022). \url{https://arxiv.org/abs/2202.00512} } \examples{ -\dontrun{ -# Perform a denoising step -result <- ddim_scheduler_step( - model_output = model_output, - timestep = timestep, - sample = sample, - eta = 0, # Deterministic sampling - prediction_type = "epsilon") +if (torch::torch_is_installed()) { + scheduler <- ddim_scheduler_create(num_inference_steps = 5) + sample <- torch::torch_randn(c(1, 4, 8, 8)) + model_output <- torch::torch_randn(c(1, 4, 8, 8)) + result <- ddim_scheduler_step( + model_output = model_output, + timestep = scheduler$timesteps[1], + sample = sample, + schedule = scheduler, + eta = 0, # Deterministic sampling + prediction_type = "epsilon") + result$shape } } diff --git a/man/dot-ltx23_pin_component.Rd b/man/dot-pin_component.Rd similarity index 50% rename from man/dot-ltx23_pin_component.Rd rename to man/dot-pin_component.Rd index 048aa48..4b22be3 100644 --- a/man/dot-ltx23_pin_component.Rd +++ b/man/dot-pin_component.Rd @@ -1,12 +1,18 @@ % tinyrox says don't edit this manually, but it can't stop you! -\name{.ltx23_pin_component} -\alias{.ltx23_pin_component} +\name{.pin_component} +\alias{.pin_component} \title{Pin a component's tensors for fast phase transfer} \usage{ -.ltx23_pin_component(module) +.pin_component(module, extra = NULL) } \arguments{ \item{module}{An nn_module on the CPU.} + +\item{extra}{Optional list of additional plain-field tensors to pin +alongside the module's parameters and buffers (e.g. an fp8 +linear's \code{weight_fp8}/\code{weight_scale} fields, which live +outside \code{parameters}/\code{buffers}). \code{set_data} mutates +each tensor in place, so the field reference stays valid.} } \value{ A list of \code{list(live, pinned)} tensor pairs, or NULL diff --git a/man/dot-ltx23_staged_offload.Rd b/man/dot-staged_offload.Rd similarity index 75% rename from man/dot-ltx23_staged_offload.Rd rename to man/dot-staged_offload.Rd index 9143965..50bc09a 100644 --- a/man/dot-ltx23_staged_offload.Rd +++ b/man/dot-staged_offload.Rd @@ -1,9 +1,9 @@ % tinyrox says don't edit this manually, but it can't stop you! -\name{.ltx23_staged_offload} -\alias{.ltx23_staged_offload} +\name{.staged_offload} +\alias{.staged_offload} \title{Return a pinned component to the CPU} \usage{ -.ltx23_staged_offload(staging) +.staged_offload(staging) } \description{ Weights are immutable during inference, so the pinned host copies diff --git a/man/dot-ltx23_staged_onload.Rd b/man/dot-staged_onload.Rd similarity index 75% rename from man/dot-ltx23_staged_onload.Rd rename to man/dot-staged_onload.Rd index 222e554..b9bb780 100644 --- a/man/dot-ltx23_staged_onload.Rd +++ b/man/dot-staged_onload.Rd @@ -1,9 +1,9 @@ % tinyrox says don't edit this manually, but it can't stop you! -\name{.ltx23_staged_onload} -\alias{.ltx23_staged_onload} +\name{.staged_onload} +\alias{.staged_onload} \title{Move a pinned component onto the compute device} \usage{ -.ltx23_staged_onload(staging, device) +.staged_onload(staging, device) } \description{ Non-blocking copies from pinned memory share the default stream, diff --git a/man/flowmatch_scheduler_create.Rd b/man/flowmatch_scheduler_create.Rd index 0a2d899..35fe1fc 100644 --- a/man/flowmatch_scheduler_create.Rd +++ b/man/flowmatch_scheduler_create.Rd @@ -66,14 +66,14 @@ Lipman, Y., Chen, R. T. Q., Ben-Hamu, H., Nickel, M., & Le, M. (2022). \url{https://arxiv.org/abs/2210.02747} } \examples{ -\dontrun{ -# Create a FlowMatch scheduler -scheduler <- flowmatch_scheduler_create( - num_train_timesteps = 1000, - shift = 1.0 -) +if (torch::torch_is_installed()) { + scheduler <- flowmatch_scheduler_create( + num_train_timesteps = 1000, + shift = 1.0 + ) -# Set timesteps for inference -scheduler <- flowmatch_set_timesteps(scheduler, num_inference_steps = 8) + # Set timesteps for inference + scheduler <- flowmatch_set_timesteps(scheduler, num_inference_steps = 8) + scheduler$timesteps } } diff --git a/man/flux2_feed_forward.Rd b/man/flux2_feed_forward.Rd index 4be65a3..05efabc 100644 --- a/man/flux2_feed_forward.Rd +++ b/man/flux2_feed_forward.Rd @@ -14,6 +14,10 @@ flux2_feed_forward(dim, dim_out = NULL, mult = 3, bias = FALSE) \item{bias}{Logical.} } +\value{ +Module whose forward(x) returns the SwiGLU-gated projection + of \code{x}, a tensor with the last axis of width \code{dim_out}. +} \description{ \code{linear_in} projects to twice the inner dim; SwiGLU gates the first half with SiLU and multiplies by the second half; diff --git a/man/flux2_load_pipeline.Rd b/man/flux2_load_pipeline.Rd index 0fbc531..e6ebefc 100644 --- a/man/flux2_load_pipeline.Rd +++ b/man/flux2_load_pipeline.Rd @@ -5,7 +5,8 @@ \usage{ flux2_load_pipeline(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4"), text_device = NULL, - attn_chunk = NULL, phase_offload = TRUE, verbose = TRUE) + attn_chunk = NULL, phase_offload = TRUE, pin = NULL, + verbose = TRUE) } \arguments{ \item{model_dir}{Quantized artifact directory (default: the @@ -24,6 +25,11 @@ fp8 when safetensors supports float8, else nf4), "fp8", or "nf4".} \item{phase_offload}{Logical. One GPU tenant per phase.} +\item{pin}{Logical or NULL. Page-lock the phase-swapped weights for +DMA-rate transfer (see \code{\link{staging}}). NULL (default) +resolves via \code{options(diffuseR.pin_staging)} then the +host-RAM-aware \code{\link{recommend}} decision.} + \item{verbose}{Logical.} } \value{ diff --git a/man/flux2_modulation.Rd b/man/flux2_modulation.Rd index 620120f..71ff778 100644 --- a/man/flux2_modulation.Rd +++ b/man/flux2_modulation.Rd @@ -12,6 +12,11 @@ flux2_modulation(dim, mod_param_sets = 2L, bias = FALSE) \item{bias}{Logical.} } +\value{ +Module whose forward(temb) returns the modulation tensor + \code{linear(silu(temb))}, holding \code{mod_param_sets} triples of + (shift, scale, gate) along the last axis. +} \description{ \code{linear(silu(temb))} producing \code{mod_param_sets} triples of (shift, scale, gate). Computed once per forward at model level and diff --git a/man/flux2_parallel_self_attention.Rd b/man/flux2_parallel_self_attention.Rd index ee36dcc..04c2598 100644 --- a/man/flux2_parallel_self_attention.Rd +++ b/man/flux2_parallel_self_attention.Rd @@ -19,6 +19,12 @@ flux2_parallel_self_attention(query_dim, heads, dim_head, mlp_ratio = 3, \item{bias}{Logical.} } +\value{ +Module whose forward(hidden_states, image_rotary_emb, + chunk_size) returns the block output [B, S, query_dim]: attention + and MLP branches computed in parallel from one fused projection, + concatenated, and projected back by a second fused layer. +} \description{ ViT-22B-style parallel block internals: one fused projection produces QKV and the SwiGLU MLP input; one fused projection consumes diff --git a/man/flux_ada_layer_norm_continuous.Rd b/man/flux_ada_layer_norm_continuous.Rd index d9e5bfd..cb8cd51 100644 --- a/man/flux_ada_layer_norm_continuous.Rd +++ b/man/flux_ada_layer_norm_continuous.Rd @@ -13,6 +13,11 @@ flux_ada_layer_norm_continuous(dim, cond_dim = dim, bias = TRUE) \item{bias}{Logical. Bias on the projection (TRUE for FLUX.1, FALSE for FLUX.2).} } +\value{ +Module whose forward(x, cond) returns \code{x} normalized and + then scaled and shifted by the conditioning embedding, a tensor of + the same shape as \code{x}. +} \description{ Scale/shift conditioning of the final norm. Note the chunk order: scale first, then shift (the reverse of adaLN-Zero). Reference: diff --git a/man/flux_attention.Rd b/man/flux_attention.Rd index 949bb69..1a3ee37 100644 --- a/man/flux_attention.Rd +++ b/man/flux_attention.Rd @@ -22,6 +22,13 @@ flux_attention(query_dim, heads, dim_head, added_kv = FALSE, pre_only = FALSE, \item{bias}{Logical. Bias on the linear projections (TRUE for FLUX.1, FALSE for FLUX.2).} } +\value{ +Module whose forward(hidden_states, encoder_hidden_states, + image_rotary_emb, chunk_size) returns the attended image stream + [B, S, query_dim]. When \code{encoder_hidden_states} is supplied + (double-stream blocks) it returns \code{list(image, text)} instead, + each projected by its own output layer. +} \description{ Multi-head attention with per-head RMS q/k norms and rotary position embeddings. With \code{added_kv = TRUE} (double-stream blocks) the diff --git a/man/flux_load_pipeline.Rd b/man/flux_load_pipeline.Rd index cef84a9..1d033fa 100644 --- a/man/flux_load_pipeline.Rd +++ b/man/flux_load_pipeline.Rd @@ -4,8 +4,8 @@ \title{Load the FLUX.1-schnell pipeline} \usage{ flux_load_pipeline(model_dir = NULL, device = "cuda", precision = NULL, - text_device = "cpu", attn_chunk = NULL, - phase_offload = TRUE, verbose = TRUE) + text_device = NULL, attn_chunk = NULL, phase_offload = TRUE, + pin = NULL, verbose = TRUE) } \arguments{ \item{model_dir}{Quantized artifact directory (default: the @@ -17,13 +17,21 @@ diffusers transformer directory for full-precision loading.} \item{precision}{"nf4" or "fp8"; NULL picks the \code{\link{flux_memory_profile}} recommendation.} -\item{text_device}{Device for the text encoders ("cpu" default; the -T5-XXL runs float32 there).} +\item{text_device}{Where the text encoders compute. NULL (default) +takes the \code{\link{flux_memory_profile}} recommendation: "cuda" +on 14 GB+ cards (T5-XXL onloads in bfloat16 for its encode) and +"cpu" below that (T5-XXL runs float32 in place, since its ~9.8 GB +GPU encode phase would not fit).} \item{attn_chunk}{Integer or NULL. Attention query-chunk override.} \item{phase_offload}{Logical. One GPU tenant per phase.} +\item{pin}{Logical or NULL. Page-lock the phase-swapped weights for +DMA-rate transfer (see \code{\link{staging}}). NULL (default) +resolves via \code{options(diffuseR.pin_staging)} then the +host-RAM-aware \code{\link{recommend}} decision.} + \item{verbose}{Logical.} } \value{ diff --git a/man/gemma3_text_model.Rd b/man/gemma3_text_model.Rd index 79e9b1d..f161f1e 100644 --- a/man/gemma3_text_model.Rd +++ b/man/gemma3_text_model.Rd @@ -8,6 +8,11 @@ gemma3_text_model(config) \arguments{ \item{config}{Model configuration list.} } +\value{ +Module whose forward(input_ids, ...) returns + \code{list(last_hidden_state, hidden_states)}: the final hidden + state [B, S, hidden_size] and the list of per-layer hidden states. +} \description{ Full Gemma3 text encoder model. } diff --git a/man/is_blackwell_gpu.Rd b/man/is_blackwell_gpu.Rd index a356f97..210de51 100644 --- a/man/is_blackwell_gpu.Rd +++ b/man/is_blackwell_gpu.Rd @@ -12,9 +12,7 @@ Logical. TRUE if Blackwell GPU detected. Blackwell GPUs (RTX 50xx) may need special handling. } \examples{ -\dontrun{ -if (is_blackwell_gpu()) { - message("Using Blackwell-compatible settings") -} -} +# Soft probe: FALSE on any machine without a Blackwell card, and on +# machines where torch has no lantern binaries. +is_blackwell_gpu() } diff --git a/man/load_gemma3_nf4.Rd b/man/load_gemma3_nf4.Rd index 577e274..7c3208b 100644 --- a/man/load_gemma3_nf4.Rd +++ b/man/load_gemma3_nf4.Rd @@ -17,7 +17,7 @@ load_gemma3_nf4(artifact_dir, device = "cuda", dtype = "bfloat16", \item{pin}{Logical. When loading to the CPU, page-lock the weights so \code{\link{encode_with_gemma3}} can swap the model to the GPU at DMA speed per encode and back for free (see -\code{\link{staging_ltx23}}). Default follows +\code{\link{staging}}). Default follows \code{options(diffuseR.pin_staging)}.} \item{verbose}{Logical.} diff --git a/man/load_to_gpu.Rd b/man/load_to_gpu.Rd index b21287a..7fb40f0 100644 --- a/man/load_to_gpu.Rd +++ b/man/load_to_gpu.Rd @@ -17,9 +17,9 @@ The module (modified in place). Moves a torch module and all its parameters to CUDA. } \examples{ -\dontrun{ -load_to_gpu(model) -output <- model(x) -offload_to_cpu(model) +if (torch::torch_is_installed()) { + model <- torch::nn_linear(4, 2) + # "cuda" needs a GPU; "cpu" is the portable round trip. + load_to_gpu(model, device = "cpu") } } diff --git a/man/ltx23_antialias_act1d.Rd b/man/ltx23_antialias_act1d.Rd index de5f9af..00e0927 100644 --- a/man/ltx23_antialias_act1d.Rd +++ b/man/ltx23_antialias_act1d.Rd @@ -10,6 +10,12 @@ ltx23_antialias_act1d(channels, ratio = 2L, kernel_size = 12L) \item{ratio,kernel_size}{Integers. Resampling config.} } +\value{ +Module whose forward(x) returns the activation applied at 2x + rate (upsample, activate, downsample), a tensor of the same shape + as \code{x}, with the aliasing the raw activation would introduce + filtered out. +} \description{ Upsample 2x, apply the activation, downsample 2x. } diff --git a/man/ltx23_attention.Rd b/man/ltx23_attention.Rd index 36dea3b..7963a98 100644 --- a/man/ltx23_attention.Rd +++ b/man/ltx23_attention.Rd @@ -28,6 +28,10 @@ ltx23_attention(query_dim, heads = 8L, kv_heads = NULL, dim_head = 64L, \item{bias,out_bias}{Logicals. Projection biases.} } +\value{ +Module whose forward(hidden_states, ...) returns the attended + states [B, S, query_dim] after the output projection. +} \description{ Attention with RMS q/k norms across heads, optional per-head output gating (LTX-2.3), separate query/key RoPE (for a2v/v2a cross diff --git a/man/ltx23_audio_causal_conv2d.Rd b/man/ltx23_audio_causal_conv2d.Rd index 0b81bc3..0a726d5 100644 --- a/man/ltx23_audio_causal_conv2d.Rd +++ b/man/ltx23_audio_causal_conv2d.Rd @@ -15,6 +15,11 @@ ltx23_audio_causal_conv2d(in_channels, out_channels, kernel_size = 3L, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(x) returns the convolved tensor, padded + so that each output frame depends only on current and earlier + input frames. +} \description{ Pads asymmetrically along the causal axis ("height" = time frames for LTX audio) before an unpadded Conv2d. diff --git a/man/ltx23_audio_decoder.Rd b/man/ltx23_audio_decoder.Rd index d29537e..0130a77 100644 --- a/man/ltx23_audio_decoder.Rd +++ b/man/ltx23_audio_decoder.Rd @@ -24,6 +24,10 @@ ltx23_audio_decoder(base_channels = 128L, output_channels = 2L, \item{mel_bins}{Integer. Output mel bins (crop/pad target).} } +\value{ +Module whose forward(x) returns the decoded mel + spectrogram reconstructed from an audio latent. +} \description{ Latents [B, 8, L, 16] -> mel spectrogram [B, 2, 4L - 3, 64]. } diff --git a/man/ltx23_audio_downsample.Rd b/man/ltx23_audio_downsample.Rd index e86a97a..9e2cda0 100644 --- a/man/ltx23_audio_downsample.Rd +++ b/man/ltx23_audio_downsample.Rd @@ -10,6 +10,10 @@ ltx23_audio_downsample(in_channels, causality_axis = "height") \item{causality_axis}{Character.} } +\value{ +Module whose forward(x) returns the strided convolution of + \code{x}, halving the downsampled axes. +} \description{ Causal zero-pad followed by a plain stride-2 conv (reference LTX2AudioDownsample; note the conv is unwrapped, so its checkpoint diff --git a/man/ltx23_audio_encoder.Rd b/man/ltx23_audio_encoder.Rd index ebeaab8..e4e9f7f 100644 --- a/man/ltx23_audio_encoder.Rd +++ b/man/ltx23_audio_encoder.Rd @@ -12,6 +12,11 @@ ltx23_audio_encoder(base_channels = 128L, in_channels = 2L, \item{base_channels,num_res_blocks,latent_channels,ch_mult,causality_axis}{See \code{\link{ltx23_audio_decoder}}.} } +\value{ +Module whose forward(x) returns the encoded audio latent, a + tensor downsampled along time and mel axes with the configured + latent channel count. +} \description{ Mel spectrogram [B, 2, T, 64] -> latent distribution moments [B, 2 * latent_channels, ceil(T/4), 16]. Structure mirrors the diff --git a/man/ltx23_audio_resnet_block.Rd b/man/ltx23_audio_resnet_block.Rd index 9b3c8a9..e7bae9f 100644 --- a/man/ltx23_audio_resnet_block.Rd +++ b/man/ltx23_audio_resnet_block.Rd @@ -11,6 +11,10 @@ ltx23_audio_resnet_block(in_channels, out_channels = NULL, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(x) returns \code{x} plus the residual + branch, a tensor of the same shape as \code{x}. +} \description{ PixelNorm -> SiLU -> causal conv, twice, with a 1x1 causal conv shortcut (\code{nin_shortcut}) on channel change. diff --git a/man/ltx23_audio_upsample.Rd b/man/ltx23_audio_upsample.Rd index 8388f7e..a63a761 100644 --- a/man/ltx23_audio_upsample.Rd +++ b/man/ltx23_audio_upsample.Rd @@ -10,6 +10,10 @@ ltx23_audio_upsample(in_channels, causality_axis = "height") \item{causality_axis}{Character.} } +\value{ +Module whose forward(x) returns the tensor upsampled 2x by + nearest-neighbour interpolation and convolved. +} \description{ Nearest 2x interpolation, causal conv, then a crop of the first element along the causal axis. diff --git a/man/ltx23_audio_vae.Rd b/man/ltx23_audio_vae.Rd index bf4cb34..8e654d4 100644 --- a/man/ltx23_audio_vae.Rd +++ b/man/ltx23_audio_vae.Rd @@ -13,6 +13,12 @@ ltx23_audio_vae(base_channels = 128L, output_channels = 2L, \item{base_channels,output_channels,num_res_blocks,latent_channels,ch_mult,causality_axis,mel_bins}{See \code{\link{ltx23_audio_decoder}}.} } +\value{ +Module bundling the audio encoder and decoder. Its + forward(z) is \code{decode(z)}, returning the mel spectrogram for a + latent; \code{$encode()} and \code{$decode()} are callable + separately. +} \description{ Encoder + decoder plus the per-channel latent statistics loaded from the checkpoint. Encoding is used for audio-conditioned generation diff --git a/man/ltx23_causal_conv3d.Rd b/man/ltx23_causal_conv3d.Rd index e7d74a0..c4ed8fc 100644 --- a/man/ltx23_causal_conv3d.Rd +++ b/man/ltx23_causal_conv3d.Rd @@ -15,6 +15,12 @@ ltx23_causal_conv3d(in_channels, out_channels, kernel_size = 3L, stride = 1L, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(hidden_states, causal) returns the 3-D + convolution of the input. With \code{causal = TRUE} the temporal + axis is left-padded by replicating the first frame, so no output + frame sees a later input frame. +} \description{ Spatial padding is handled by the convolution; temporal padding replicates the first frame (causal) or both edge frames (non-causal), diff --git a/man/ltx23_connector_transformer_1d.Rd b/man/ltx23_connector_transformer_1d.Rd index 058e6f5..8b38659 100644 --- a/man/ltx23_connector_transformer_1d.Rd +++ b/man/ltx23_connector_transformer_1d.Rd @@ -22,6 +22,12 @@ sequence length must be divisible by it).} \item{rope_base_seq_len,rope_theta,rope_double_precision,rope_type}{RoPE config.} } +\value{ +Module whose forward(hidden_states, attention_mask, + attn_mask_binarize_threshold) returns + \code{list(hidden_states, attention_mask)}: the transformed + sequence and the (possibly binarized) mask that accompanies it. +} \description{ Replaces padded positions with learnable registers (valid tokens are front-aligned in their original order; the tail is filled with diff --git a/man/ltx23_downsample1d.Rd b/man/ltx23_downsample1d.Rd index ecc73cb..9ebe0f4 100644 --- a/man/ltx23_downsample1d.Rd +++ b/man/ltx23_downsample1d.Rd @@ -10,6 +10,10 @@ ltx23_downsample1d(ratio = 2L, kernel_size = NULL) \item{kernel_size}{Integer or NULL (default 6*ratio rounded even).} } +\value{ +Module whose forward(x) returns \code{x} low-pass filtered + and decimated by \code{ratio} along the time axis. +} \description{ Anti-aliasing 1D downsampler (low-pass then stride) } diff --git a/man/ltx23_feed_forward.Rd b/man/ltx23_feed_forward.Rd index e142ead..aaa197f 100644 --- a/man/ltx23_feed_forward.Rd +++ b/man/ltx23_feed_forward.Rd @@ -10,6 +10,11 @@ ltx23_feed_forward(dim, mult = 4L) \item{mult}{Integer. Hidden dimension multiplier.} } +\value{ +Module whose forward(x) returns the projected states passed + through a tanh-approximated GELU, a tensor of the same shape as + \code{x}. +} \description{ Linear -> GELU (tanh approximation) -> Linear with 4x hidden dim, matching diffusers \code{FeedForward(activation_fn="gelu-approximate")} diff --git a/man/ltx23_fp8_linear.Rd b/man/ltx23_fp8_linear.Rd index dd87539..c5c3509 100644 --- a/man/ltx23_fp8_linear.Rd +++ b/man/ltx23_fp8_linear.Rd @@ -10,6 +10,13 @@ ltx23_fp8_linear(out_features, in_features, bias = TRUE) \item{out_features,in_features}{Integers.} } +\value{ +Module whose forward(x) returns the linear projection of + \code{x}, with the fp8 weight bytes transferred and cast up to the + compute dtype for the matmul. Same result as an + \code{nn_linear} of the same shape, at a quarter of the resident + weight bytes. +} \description{ Weight lives as float8_e4m3fn plus a float32 scale in plain module fields (so \code{$to(device)} moves only the bias); the forward pass diff --git a/man/ltx23_latent_upsampler.Rd b/man/ltx23_latent_upsampler.Rd index fa4cae1..c7737db 100644 --- a/man/ltx23_latent_upsampler.Rd +++ b/man/ltx23_latent_upsampler.Rd @@ -13,6 +13,11 @@ ltx23_latent_upsampler(in_channels = 128L, mid_channels = 1024L, \item{num_blocks_per_stage}{Integer.} } +\value{ +Module whose forward(hidden_states) returns the 2x + spatially upscaled latent, a tensor with the same batch, channel + and frame counts and doubled height and width. +} \description{ Latents [B, 128, F, H, W] -> [B, 128, F, 2H, 2W]. } diff --git a/man/ltx23_mel_stft.Rd b/man/ltx23_mel_stft.Rd index 5d36d8e..74cf31c 100644 --- a/man/ltx23_mel_stft.Rd +++ b/man/ltx23_mel_stft.Rd @@ -9,6 +9,10 @@ ltx23_mel_stft(filter_length = 512L, hop_length = 80L, window_length = 512L, \arguments{ \item{filter_length,hop_length,window_length,num_mel_channels}{Integers.} } +\value{ +Module whose forward(waveform) returns the log-mel + spectrogram [B, n_mels, frames], clamped at 1e-5 before the log. +} \description{ Causal log-mel spectrogram with checkpoint-loaded bases } diff --git a/man/ltx23_nf4_linear.Rd b/man/ltx23_nf4_linear.Rd index 2bc15c6..5896af0 100644 --- a/man/ltx23_nf4_linear.Rd +++ b/man/ltx23_nf4_linear.Rd @@ -10,6 +10,12 @@ ltx23_nf4_linear(out_features, in_features, bias = TRUE) \item{out_features,in_features}{Integers.} } +\value{ +Module whose forward(x) returns the linear projection of + \code{x}, dequantizing the NF4 weight into a reusable buffer first. + Same result as an \code{nn_linear} of the same shape, at roughly an + eighth of the resident weight bytes. +} \description{ Packed weights and per-block scales are registered as buffers, so they move with the module (uint8 packs are untouched by dtype diff --git a/man/ltx23_per_channel_rms_norm.Rd b/man/ltx23_per_channel_rms_norm.Rd index 74616f7..eb0e677 100644 --- a/man/ltx23_per_channel_rms_norm.Rd +++ b/man/ltx23_per_channel_rms_norm.Rd @@ -8,6 +8,10 @@ ltx23_per_channel_rms_norm(eps = 1e-08) \arguments{ \item{eps}{Numeric. Stability epsilon.} } +\value{ +Module whose forward(x) returns \code{x} divided by its + per-channel root mean square, a tensor of the same shape. +} \description{ Normalizes by the root-mean-square across the channel dimension (dim 2 of [B, C, F, H, W]); no learned parameters. diff --git a/man/ltx23_rms_norm.Rd b/man/ltx23_rms_norm.Rd index b7214ae..93f04b3 100644 --- a/man/ltx23_rms_norm.Rd +++ b/man/ltx23_rms_norm.Rd @@ -12,6 +12,11 @@ ltx23_rms_norm(dim, eps = 1e-06, elementwise_affine = TRUE) \item{elementwise_affine}{Logical. Learn a scale weight.} } +\value{ +Module whose forward(x) returns \code{x} RMS-normalized over + the last axis and cast back to the input dtype, a tensor of the + same shape. +} \description{ Variance is computed in float32; the result is cast back to the input dtype (or the weight dtype when elementwise affine). diff --git a/man/ltx23_rotary_pos_embed.Rd b/man/ltx23_rotary_pos_embed.Rd index 8d6dcdd..ad8ff38 100644 --- a/man/ltx23_rotary_pos_embed.Rd +++ b/man/ltx23_rotary_pos_embed.Rd @@ -38,6 +38,11 @@ coordinates are normalized against.} \item{sampling_rate,hop_length}{Integers. Audio spectrogram params.} } +\value{ +Module whose forward(coords, device) returns + \code{list(cos_freqs, sin_freqs)}, the two rotary tables to apply + to queries and keys. +} \description{ Computes RoPE cos/sin frequency tensors from spatiotemporal patch coordinates. Video coordinates are 3D (frames scaled to seconds via diff --git a/man/ltx23_rotary_pos_embed_1d.Rd b/man/ltx23_rotary_pos_embed_1d.Rd index 9324bd9..92a9a63 100644 --- a/man/ltx23_rotary_pos_embed_1d.Rd +++ b/man/ltx23_rotary_pos_embed_1d.Rd @@ -20,6 +20,11 @@ ltx23_rotary_pos_embed_1d(dim, base_seq_len = 4096L, theta = 10000, \item{num_attention_heads}{Integer. For the split per-head layout.} } +\value{ +Module whose forward(batch_size, pos, device) returns + \code{list(cos_freqs, sin_freqs)}, the 1-D rotary tables for a + sequence of length \code{pos}. +} \description{ 1D rotary embeddings for the text connectors } diff --git a/man/ltx23_snake_beta.Rd b/man/ltx23_snake_beta.Rd index ec90b66..59f7cd7 100644 --- a/man/ltx23_snake_beta.Rd +++ b/man/ltx23_snake_beta.Rd @@ -10,6 +10,11 @@ ltx23_snake_beta(channels, eps = 1e-09) \item{eps}{Numeric.} } +\value{ +Module whose forward(hidden_states) returns the Snake + activation \code{x + sin(alpha * x)^2 / beta}, a tensor of the same + shape as the input. +} \description{ \code{x + (1 / (beta + eps)) * sin(x * alpha)^2} with per-channel log-scale alpha/beta parameters. diff --git a/man/ltx23_text_connectors.Rd b/man/ltx23_text_connectors.Rd index ab6722c..73e21e1 100644 --- a/man/ltx23_text_connectors.Rd +++ b/man/ltx23_text_connectors.Rd @@ -45,6 +45,13 @@ Gemma3-12B).} \item{video_hidden_dim,audio_hidden_dim}{Integers. Projection targets (DiT inner dims: 4096 / 2048).} } +\value{ +Module whose forward(text_encoder_hidden_states, + attention_mask) returns \code{list(video_text_embedding, + audio_text_embedding, attention_mask)}: the caption states adapted + for the video and audio cross-attention streams, plus the binary + mask to use with them. +} \description{ Takes raw stacked per-layer text encoder hidden states and produces the video and audio text embeddings for the DiT: per-token RMS norm, diff --git a/man/ltx23_transformer.Rd b/man/ltx23_transformer.Rd index 5a1837f..2dcc039 100644 --- a/man/ltx23_transformer.Rd +++ b/man/ltx23_transformer.Rd @@ -58,6 +58,12 @@ the a2v/v2a gates).} \item{gated_attn,cross_attn_mod,audio_gated_attn,audio_cross_attn_mod,perturbed_attn}{LTX-2.3 feature flags (all TRUE for the 2.3 checkpoints).} } +\value{ +Module whose forward(hidden_states, ...) returns + \code{list(sample, audio_sample)}: the predicted velocity for the + video latent tokens and, when the audio branch is active, for the + audio latent tokens (\code{audio_sample} is NULL otherwise). +} \description{ Dual-stream audio/video DiT. Text embeddings arrive already projected to the video (\code{inner_dim}) and audio (\code{audio_inner_dim}) diff --git a/man/ltx23_transformer_block.Rd b/man/ltx23_transformer_block.Rd index 6391570..bf4c04f 100644 --- a/man/ltx23_transformer_block.Rd +++ b/man/ltx23_transformer_block.Rd @@ -36,6 +36,12 @@ ltx23_transformer_block(dim, num_attention_heads, attention_head_dim, \item{video_cross_attn_adaln,audio_cross_attn_adaln}{Logicals. LTX-2.3 text cross-attention modulation (9 mod params instead of 6).} } +\value{ +Module whose forward(hidden_states, ...) returns + \code{list(hidden_states, audio_hidden_states)}, the video and audio + streams after self-attention, cross-attention and the feed-forward, + each the same shape as its input. +} \description{ Dual-stream (video + audio) block: modulated self-attention per modality, text cross-attention per modality (with LTX-2.3 query and diff --git a/man/ltx23_upsample1d.Rd b/man/ltx23_upsample1d.Rd index bdf5800..af2864e 100644 --- a/man/ltx23_upsample1d.Rd +++ b/man/ltx23_upsample1d.Rd @@ -16,6 +16,11 @@ ltx23_upsample1d(ratio = 2L, kernel_size = NULL, window_type = "kaiser", \item{persistent}{Logical. Register the filter as a buffer (present in checkpoints); FALSE stores the computed filter as a plain field.} } +\value{ +Module whose forward(x) returns \code{x} interpolated up by + \code{ratio} along the time axis, with the filter padding trimmed + off. +} \description{ Anti-aliasing 1D upsampler (transposed low-pass) } diff --git a/man/ltx23_video_decoder3d.Rd b/man/ltx23_video_decoder3d.Rd index 818ee26..39204dc 100644 --- a/man/ltx23_video_decoder3d.Rd +++ b/man/ltx23_video_decoder3d.Rd @@ -37,6 +37,11 @@ the mid block after reversal).} \item{patch_size,patch_size_t}{Integers.} } +\value{ +Module whose forward(hidden_states, causal) returns the + decoded pixel tensor [B, 3, F, H, W], with the final patch axes + flattened back into height and width. +} \description{ Latents [B, 128, F, H, W] -> pixel video [B, 3, 8F - 7, 32H, 32W]. Block channel lists are given encoder-side (as in the config) and diff --git a/man/ltx23_video_down_block3d.Rd b/man/ltx23_video_down_block3d.Rd index e8004be..1841b4e 100644 --- a/man/ltx23_video_down_block3d.Rd +++ b/man/ltx23_video_down_block3d.Rd @@ -21,6 +21,11 @@ ltx23_video_down_block3d(in_channels, out_channels = NULL, num_layers = 1L, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(hidden_states, causal) returns the + stage output: the resnet stack applied in sequence, then the + optional downsampler. +} \description{ ResNet stack (at the input channel count) followed by a pixel-unshuffle downsampler that also changes the channel count. diff --git a/man/ltx23_video_downsampler3d.Rd b/man/ltx23_video_downsampler3d.Rd index 0afebb6..3203165 100644 --- a/man/ltx23_video_downsampler3d.Rd +++ b/man/ltx23_video_downsampler3d.Rd @@ -13,6 +13,11 @@ ltx23_video_downsampler3d(in_channels, out_channels, stride = c(1L, 1L, 1L), \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(hidden_states, causal) returns the + space-to-depth downsampled tensor plus its residual: spatial and + temporal extents shrink by \code{stride}, channels grow to match. +} \description{ Conv followed by space/time-to-channel rearrangement, plus a grouped channel-mean residual of the same rearrangement. diff --git a/man/ltx23_video_encoder3d.Rd b/man/ltx23_video_encoder3d.Rd index e3ddacc..ee9702a 100644 --- a/man/ltx23_video_encoder3d.Rd +++ b/man/ltx23_video_encoder3d.Rd @@ -30,6 +30,11 @@ ltx23_video_encoder3d(in_channels = 3L, out_channels = 128L, \item{patch_size,patch_size_t}{Integers. Pixel patchification.} } +\value{ +Module whose forward(hidden_states, causal) returns the + encoded video latent [B, C, F, H, W], with the last channel + repeated to carry the per-channel scale expected downstream. +} \description{ Pixel video [B, 3, F, H, W] -> latent statistics [B, 2 * latent_channels, F/8, H/32, W/32] (mean and a uniform log-var diff --git a/man/ltx23_video_mid_block3d.Rd b/man/ltx23_video_mid_block3d.Rd index ac45a31..dba8e4c 100644 --- a/man/ltx23_video_mid_block3d.Rd +++ b/man/ltx23_video_mid_block3d.Rd @@ -15,6 +15,10 @@ ltx23_video_mid_block3d(in_channels, num_layers = 1L, resnet_eps = 1e-06, \item{spatial_padding_mode}{Character.} } +\value{ +Module whose forward(hidden_states, causal) returns the + bottleneck output, a tensor of the same shape as the input. +} \description{ A plain ResNet stack at a fixed channel count. } diff --git a/man/ltx23_video_resnet_block3d.Rd b/man/ltx23_video_resnet_block3d.Rd index 2506593..bdb7636 100644 --- a/man/ltx23_video_resnet_block3d.Rd +++ b/man/ltx23_video_resnet_block3d.Rd @@ -13,6 +13,10 @@ ltx23_video_resnet_block3d(in_channels, out_channels = NULL, eps = 1e-06, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(inputs, causal) returns \code{inputs} + plus the residual branch, a tensor of the same shape. +} \description{ PerChannelRMSNorm -> SiLU -> causal conv, twice, with a LayerNorm + 1x1 Conv3d shortcut when the channel count changes. diff --git a/man/ltx23_video_up_block3d.Rd b/man/ltx23_video_up_block3d.Rd index f234c16..2cc16af 100644 --- a/man/ltx23_video_up_block3d.Rd +++ b/man/ltx23_video_up_block3d.Rd @@ -26,6 +26,11 @@ ltx23_video_up_block3d(in_channels, out_channels = NULL, num_layers = 1L, \item{in_channels,out_channels}{Integers.} } +\value{ +Module whose forward(hidden_states, causal) returns the + stage output: the optional input projection and upsampler, then the + resnet stack applied in sequence. +} \description{ Optional channel-changing conv-in ResNet, pixel-shuffle upsampler, then a ResNet stack at the output channel count. diff --git a/man/ltx23_video_upsampler3d.Rd b/man/ltx23_video_upsampler3d.Rd index e2d25ca..c060011 100644 --- a/man/ltx23_video_upsampler3d.Rd +++ b/man/ltx23_video_upsampler3d.Rd @@ -17,6 +17,11 @@ ltx23_video_upsampler3d(in_channels, stride = c(1L, 1L, 1L), residual = FALSE, \item{spatial_padding_mode}{Character.} } +\value{ +Module whose forward(hidden_states, causal) returns the + depth-to-space upsampled tensor: spatial and temporal extents grow + by \code{stride}, channels shrink to match. +} \description{ Conv followed by channel-to-space/time rearrangement, with an optional channel-repeat residual and an upscale factor that divides the conv diff --git a/man/ltx23_video_vae.Rd b/man/ltx23_video_vae.Rd index 6f4d047..6f673d2 100644 --- a/man/ltx23_video_vae.Rd +++ b/man/ltx23_video_vae.Rd @@ -35,6 +35,11 @@ ltx23_video_vae(in_channels = 3L, out_channels = 3L, latent_channels = 128L, \item{encoder_spatial_padding_mode,decoder_spatial_padding_mode}{Characters.} } +\value{ +Module bundling the video encoder and decoder. Its + forward(z) is \code{decode(z)}, returning pixels for a latent; + \code{$encode()} and \code{$decode()} are callable separately. +} \description{ Encoder + decoder + per-channel latent statistics (loaded from the checkpoint's \code{per_channel_statistics}). The checkpoint's diff --git a/man/ltx23_vocoder.Rd b/man/ltx23_vocoder.Rd index e9ecaa1..78dfcbb 100644 --- a/man/ltx23_vocoder.Rd +++ b/man/ltx23_vocoder.Rd @@ -28,6 +28,10 @@ ltx23_vocoder(in_channels = 128L, hidden_channels = 1536L, out_channels = 2L, \item{antialias_ratio,antialias_kernel_size}{Integers.} } +\value{ +Module whose forward(hidden_states, time_last) returns the + synthesized waveform [B, 1, samples] for a mel spectrogram. +} \description{ Mel spectrogram [B, C, T, M] -> waveform [B, out_channels, samples]. Channel and mel dims are flattened into conv channels; each upsample diff --git a/man/ltx23_vocoder_resblock.Rd b/man/ltx23_vocoder_resblock.Rd index 8d5f7db..1d6abff 100644 --- a/man/ltx23_vocoder_resblock.Rd +++ b/man/ltx23_vocoder_resblock.Rd @@ -15,6 +15,11 @@ ltx23_vocoder_resblock(channels, kernel_size = 3L, dilations = c(1L, 3L, 5L), \item{antialias_ratio,antialias_kernel_size}{Integers.} } +\value{ +Module whose forward(x) returns \code{x} after the dilated + convolution pairs have been added back as residuals, a tensor of + the same shape. +} \description{ Dilated conv pairs, each preceded by an anti-aliased SnakeBeta activation, with residual connections. diff --git a/man/ltx23_vocoder_with_bwe.Rd b/man/ltx23_vocoder_with_bwe.Rd index 4cb2f6b..b5c8f93 100644 --- a/man/ltx23_vocoder_with_bwe.Rd +++ b/man/ltx23_vocoder_with_bwe.Rd @@ -38,6 +38,11 @@ re-analysis configuration.} \item{input_sampling_rate,output_sampling_rate}{Integers.} } +\value{ +Module whose forward(mel_spec) returns the + bandwidth-extended waveform [B, 1, samples], trimmed to the sample + count implied by the input frames and the rate ratio. +} \description{ Full mel [B, 2, T, 64] -> 48 kHz stereo waveform pipeline: 16 kHz vocoder, causal mel re-analysis, BWE vocoder residual added to a diff --git a/man/memory_flux.Rd b/man/memory_flux.Rd index 083121d..cbe97da 100644 --- a/man/memory_flux.Rd +++ b/man/memory_flux.Rd @@ -4,6 +4,8 @@ \title{FLUX Memory Profiles} \description{ 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. +12B transformer runs NF4 (~7 GB) or fp8 (~12 GB), phase-onloaded to +the GPU for denoise; the T5-XXL text encoder phase-onloads to the GPU +(bfloat16, pinned) on 14 GB+ cards and computes on the CPU (float32) +below that, where its ~9.8 GB encode phase does not fit. } diff --git a/man/offload_to_cpu.Rd b/man/offload_to_cpu.Rd index 5cbdd33..94e1fd3 100644 --- a/man/offload_to_cpu.Rd +++ b/man/offload_to_cpu.Rd @@ -17,9 +17,8 @@ The module (modified in place). Moves a torch module and all its parameters to CPU. } \examples{ -\dontrun{ -model$to(device = "cuda") -output <- model(x) -offload_to_cpu(model) +if (torch::torch_is_installed()) { + model <- torch::nn_linear(4, 2) + offload_to_cpu(model) } } diff --git a/man/print.bpe_tokenizer.Rd b/man/print.bpe_tokenizer.Rd index 91eb2cd..b4e84d3 100644 --- a/man/print.bpe_tokenizer.Rd +++ b/man/print.bpe_tokenizer.Rd @@ -10,6 +10,10 @@ \item{...}{Additional arguments (ignored).} } +\value{ +Invisibly returns \code{x}. Called for the side effect of + printing a summary of the tokenizer to the console. +} \description{ Print BPE Tokenizer } diff --git a/man/recommend.Rd b/man/recommend.Rd index d24b3d9..74c3963 100644 --- a/man/recommend.Rd +++ b/man/recommend.Rd @@ -57,29 +57,32 @@ pipeline uses \code{\link{ltx23_memory_profile}} for frame-aware placement. The pinning decision: phase-swapped weights are page-locked host -copies (see \code{\link{staging_ltx23}}) that transfer at DMA rate - +copies (see \code{\link{staging}}) that transfer at DMA rate - but pinned pages are unswappable, so on small-RAM machines they turn memory pressure into OOM kills. \code{pin} is TRUE when available host RAM covers the model's pinned set twice over, FALSE below that, FALSE on the cpu tier (nothing stages), and TRUE when RAM cannot be -detected (page-locking already fails soft per component). Today the -LTX pipeline and Gemma3 encoder loaders take \code{pin} arguments -and stage pinned weights (see \code{\link{staging_ltx23}}); the -image-model loaders do not stage weights yet, so for them \code{pin} -is forward-looking policy with no consumer. -\code{options(diffuseR.pin_staging)} is the global switch. +detected (page-locking already fails soft per component). The LTX +pipeline, the Gemma3 encoder, and the FLUX-family image loaders +(flux1, flux2, zimage) take \code{pin} arguments and stage pinned +weights (see \code{\link{staging}}); the SD-family loaders place +components statically and do not phase-swap, so \code{pin} is inert +for them. \code{options(diffuseR.pin_staging)} is the global switch. } \examples{ -\dontrun{ -# Auto-detect VRAM and probe the installed safetensors -recommend("flux2") - -# A 16 GB card without float8 support: fp8 wanted, nf4 recommended +# Stating vram_gb and st_caps makes the policy deterministic: no GPU +# and no installed safetensors needed. r <- recommend("flux1", vram_gb = 16, st_caps = list(bfloat16 = TRUE, float8_e4m3fn = FALSE)) -r$precision # "nf4" +r$precision # "nf4": fp8 fits the card, but cannot be read r$fork_suggested # TRUE cat(r$note) # the fork-or-nf4 message -} + +# Same card, once safetensors can read float8 +recommend("flux1", vram_gb = 16, + st_caps = list(bfloat16 = TRUE, float8_e4m3fn = TRUE))$precision + +# Auto-detect VRAM and probe the installed safetensors +str(recommend("flux2")) } diff --git a/man/save_image.Rd b/man/save_image.Rd index e9e6365..8f3fe70 100644 --- a/man/save_image.Rd +++ b/man/save_image.Rd @@ -20,7 +20,9 @@ Converts a Torch tensor to a normalized RGB image array, saves it as a PNG file, and optionally displays it in the RStudio Viewer pane using `grid::grid.raster()`. } \examples{ -\dontrun{ -save_image(output_tensor, "sample.png") -} +img <- array(runif(32 * 32 * 3), dim = c(32, 32, 3)) +out <- file.path(tempdir(), "sample.png") +save_image(img, out) +file.exists(out) +unlink(out) } diff --git a/man/save_video.Rd b/man/save_video.Rd index 8608d49..e14d3c2 100644 --- a/man/save_video.Rd +++ b/man/save_video.Rd @@ -30,14 +30,19 @@ Invisibly returns the output file path. Saves a video array to a file in various formats. } \examples{ -\dontrun{ -# Save as MP4 -save_video(video_array, "output.mp4", fps = 24) - -# Save as GIF -save_video(video_array, "output.gif", fps = 10) - -# Save as individual frames -save_video(video_array, "frames/", format = "frames") -} +video <- array(runif(4 * 16 * 16 * 3), dim = c(4, 16, 16, 3)) + +# Individual PNG frames need no external encoder. +frame_dir <- file.path(tempdir(), "frames") +save_video(video, frame_dir, format = "frames", verbose = FALSE) +length(list.files(frame_dir, pattern = "[.]png$")) +unlink(frame_dir, recursive = TRUE) + +# MP4 and GIF need an ffmpeg binary or the 'av' package. Not shown as +# a live example: both backends hand off to an ffmpeg process that +# inherits this session's stdin, and R CMD check feeds the example +# script to R on stdin, so the encoder eats a byte of the script and +# every later example parses one character short. +# save_video(video, "output.mp4", fps = 24) +# save_video(video, "output.gif", fps = 10) } diff --git a/man/scheduler_add_noise.Rd b/man/scheduler_add_noise.Rd index 72376bd..f6c0fc8 100644 --- a/man/scheduler_add_noise.Rd +++ b/man/scheduler_add_noise.Rd @@ -35,14 +35,16 @@ Where alpha_cumprod is the cumulative product of (1-beta) values up to the specified timestep, with beta being the noise schedule. } \examples{ -\dontrun{ -# Assuming we have latents, noise, and a scheduler -noised_latents <- scheduler_add_noise( - original_latents = latents, - noise = torch::torch_randn_like(latents), - timestep = scheduler$timesteps[1], - scheduler_obj = scheduler -) +if (torch::torch_is_installed()) { + scheduler <- ddim_scheduler_create(num_inference_steps = 5) + latents <- torch::torch_randn(c(1, 4, 8, 8)) + noised_latents <- scheduler_add_noise( + original_latents = latents, + noise = torch::torch_randn_like(latents), + timestep = scheduler$timesteps[1], + scheduler_obj = scheduler + ) + noised_latents$shape } } diff --git a/man/sdxl_memory_profile.Rd b/man/sdxl_memory_profile.Rd index f008493..3a008ce 100644 --- a/man/sdxl_memory_profile.Rd +++ b/man/sdxl_memory_profile.Rd @@ -31,11 +31,11 @@ Each profile also specifies: - max_resolution: maximum image dimension } \examples{ -\dontrun{ -# Auto-detect profile -profile <- sdxl_memory_profile() +# A stated VRAM budget is deterministic and needs no GPU. +str(sdxl_memory_profile(vram_gb = 8)) -# Specific VRAM -profile <- sdxl_memory_profile(vram_gb = 8) -} +str(sdxl_memory_profile(vram_gb = 24)) + +# Auto-detect free VRAM on this machine. +str(sdxl_memory_profile()) } diff --git a/man/st_caps.Rd b/man/st_caps.Rd index 2dfbb1d..d79ec34 100644 --- a/man/st_caps.Rd +++ b/man/st_caps.Rd @@ -3,10 +3,13 @@ \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: +The CRAN build of safetensors 0.2.1 reads bfloat16 but cannot write +it, and has no float8 support at all. Both fixes merged upstream on +2026-07-31 (mlverse/safetensors#11 for bfloat16 write, #13 for +float8) without a version bump, so the installed version number +cannot tell you which build you have. That is why every gate here is +a runtime probe: write a tiny tensor, read it back, cache the answer. +Two capabilities matter and they differ: } \details{ \itemize{ diff --git a/man/staging_ltx23.Rd b/man/staging.Rd similarity index 64% rename from man/staging_ltx23.Rd rename to man/staging.Rd index d7fd3a5..ced8b5f 100644 --- a/man/staging_ltx23.Rd +++ b/man/staging.Rd @@ -1,12 +1,13 @@ % tinyrox says don't edit this manually, but it can't stop you! -\name{staging_ltx23} +\name{staging} +\alias{staging} \alias{staging_ltx23} \title{Pinned Staging for Phase-Sequential Components} \description{ Phase offloading moves each large component (transformer, -connectors, VAEs, vocoder) between CPU and GPU every render. From -pageable memory those copies run through the driver's bounce -buffer at a fraction of PCIe speed; page-locked (pinned) host +connectors, VAEs, vocoder, text encoders) between CPU and GPU every +render. From pageable memory those copies run through the driver's +bounce buffer at a fraction of PCIe speed; page-locked (pinned) host memory transfers by DMA at full rate. Each component's parameters and buffers are pinned once at load; onload swaps every tensor to a non-blocking GPU copy of its pinned source, and offload simply @@ -22,8 +23,11 @@ post byte-LUT (768x512x49, NF4, RTX 5060 Ti): ~7s saved per render identical, the delta is pure transfer), so pinning breaks even on the second render and costs a single-render session ~2s net. On by default; page-locking failure falls back silently per component, -and \code{options(diffuseR.pin_staging = FALSE)} before -\code{ltx23_load_pipeline} opts out (e.g. under host memory -pressure, where unswappable pages turn thrashing into OOM). +and \code{options(diffuseR.pin_staging = FALSE)} before the loader +opts out (e.g. under host memory pressure, where unswappable pages +turn thrashing into OOM). The LTX pipeline, the Gemma3 encoder, and +the FLUX-family image loaders (flux1, flux2, zimage) all stage +pinned weights; \code{\link{recommend}} computes the RAM-aware +\code{pin} default per model. } diff --git a/man/vae_decoder_native.Rd b/man/vae_decoder_native.Rd index c6823f8..b11c934 100644 --- a/man/vae_decoder_native.Rd +++ b/man/vae_decoder_native.Rd @@ -26,10 +26,20 @@ Native R torch implementation of the SDXL VAE decoder. Replaces TorchScript decoder for better GPU compatibility. } \examples{ +if (torch::torch_is_installed()) { + # A small decoder; the SD/SDXL defaults are far too large to build + # inside an example. + decoder <- vae_decoder_native(latent_channels = 4, + block_channels = 32, + norm_groups = 32) + latents <- torch::torch_randn(c(1, 4, 8, 8)) + image <- torch::with_no_grad(decoder(latents)) + image$shape +} + +# Real weights come from a downloaded checkpoint. \dontrun{ decoder <- vae_decoder_native() load_decoder_weights(decoder, "path/to/decoder.pt") -latents <- torch::torch_randn(c(1, 4, 64, 64)) -image <- decoder(latents) } } diff --git a/man/vram_report.Rd b/man/vram_report.Rd index 647ecff..7188224 100644 --- a/man/vram_report.Rd +++ b/man/vram_report.Rd @@ -15,7 +15,7 @@ Invisibly returns a list with used and free VRAM in GB. Prints current VRAM usage from nvidia-smi. } \examples{ -\dontrun{ -vram_report("After model load") +if (torch::torch_is_installed()) { + vram_report("After model load") } } diff --git a/man/zimage_block.Rd b/man/zimage_block.Rd index 04f67a1..e653afe 100644 --- a/man/zimage_block.Rd +++ b/man/zimage_block.Rd @@ -14,6 +14,12 @@ zimage_block(dim, n_heads, norm_eps = 1e-05, modulation = TRUE) \item{modulation}{Logical. Whether the block is timestep-modulated.} } +\value{ +Module whose forward(x, freqs, adaln_input, chunk_size) + returns the residual block output, a tensor of the same shape as + \code{x}. \code{adaln_input} is used only when the block was built + with \code{modulation = TRUE}. +} \description{ Sandwich-norm residual block shared by the noise refiner, the context refiner and the main trunk. With \code{modulation = TRUE} the block diff --git a/man/zimage_feed_forward.Rd b/man/zimage_feed_forward.Rd index 7c49d8f..88e970e 100644 --- a/man/zimage_feed_forward.Rd +++ b/man/zimage_feed_forward.Rd @@ -10,6 +10,10 @@ zimage_feed_forward(dim, hidden_dim) \item{hidden_dim}{Integer. Hidden width.} } +\value{ +Module whose forward(x) returns \code{w2(silu(w1(x)) * w3(x))}, + a tensor of the same shape as \code{x}. +} \description{ w2(silu(w1(x)) * w3(x)) with all three linears bias-free. The hidden width is int(dim / 3 * 8). diff --git a/man/zimage_final_layer.Rd b/man/zimage_final_layer.Rd index 9642f70..27af07e 100644 --- a/man/zimage_final_layer.Rd +++ b/man/zimage_final_layer.Rd @@ -11,6 +11,11 @@ zimage_final_layer(hidden_size, out_channels) \item{out_channels}{Integer. Patch output dim (patch^2 * f_patch * latent channels).} } +\value{ +Module whose forward(x, c) returns the token-to-patch + projection [B, S, out_channels], ready for unpatchifying into a + latent. +} \description{ Parameterless LayerNorm scaled by 1 + adaLN(c) (scale only, no shift), then the token-to-patch projection. diff --git a/man/zimage_load_pipeline.Rd b/man/zimage_load_pipeline.Rd index 9466ef5..b7496c8 100644 --- a/man/zimage_load_pipeline.Rd +++ b/man/zimage_load_pipeline.Rd @@ -5,7 +5,8 @@ \usage{ zimage_load_pipeline(model_dir = NULL, device = "cuda", precision = c("auto", "fp8", "nf4"), text_device = NULL, - attn_chunk = NULL, phase_offload = TRUE, verbose = TRUE) + attn_chunk = NULL, phase_offload = TRUE, pin = NULL, + verbose = TRUE) } \arguments{ \item{model_dir}{Quantized artifact directory (default: the @@ -24,6 +25,11 @@ fp8 when safetensors supports float8, else nf4), "fp8", or "nf4".} \item{phase_offload}{Logical. One GPU tenant per phase.} +\item{pin}{Logical or NULL. Page-lock the phase-swapped weights for +DMA-rate transfer (see \code{\link{staging}}). NULL (default) +resolves via \code{options(diffuseR.pin_staging)} then the +host-RAM-aware \code{\link{recommend}} decision.} + \item{verbose}{Logical.} } \value{ diff --git a/man/zimage_t_embedder.Rd b/man/zimage_t_embedder.Rd index f8c3b73..a733d12 100644 --- a/man/zimage_t_embedder.Rd +++ b/man/zimage_t_embedder.Rd @@ -12,6 +12,10 @@ zimage_t_embedder(out_size, mid_size = 1024L, freq_size = 256L) \item{freq_size}{Integer. Sinusoid width. Default 256.} } +\value{ +Module whose forward(t) returns the timestep embedding + [B, out_size]. +} \description{ 256-dim cos-first sinusoid (computed in float32) through a Linear-SiLU-Linear MLP. The model feeds t * t_scale with the diff --git a/man/zimage_transformer.Rd b/man/zimage_transformer.Rd index af81f4b..ba6281a 100644 --- a/man/zimage_transformer.Rd +++ b/man/zimage_transformer.Rd @@ -35,6 +35,13 @@ c(32, 48, 48).} \item{f_patch_size}{Integer. Temporal patch size. Default 1.} } +\value{ +Module whose forward(x, t, cap_feats, chunk_size) returns the + predicted velocity for the single latent, a tensor [C, F, H, W] + matching the shape of \code{x}. Note that the checkpoint negates + this output and consumes a reversed timestep; see + \code{\link{txt2img_zimage}}. +} \description{ Fresh R port of ZImageTransformer2DModel from the diffusers reference (Apache-2.0, src/diffusers/models/transformers/transformer_z_image.py). diff --git a/vignettes/performance-levers.md b/vignettes/performance-levers.md index 7c7d1bb..f9b1c8b 100644 --- a/vignettes/performance-levers.md +++ b/vignettes/performance-levers.md @@ -16,9 +16,10 @@ This vignette is the map. The machine-readable version of the same policy is `recommend()`, which inspects your VRAM, host RAM, and installed `safetensors` capabilities and returns a configuration. Today `flux_memory_profile()` delegates to it and `serve()` consults -it to pick between built LTX artifacts; for everything else it is -advisory — call it and pass its fields to the loaders yourself. Full -consumption (including `pin`) is arriving model by model. +it to pick between built LTX artifacts; the FLUX-family loaders resolve +their `pin` (and FLUX.1 its `text_device`) through it. For the rest of +the fields it is advisory — call it and pass them to the loaders +yourself. ## Lever 1: weight precision @@ -30,7 +31,7 @@ carries an exact census of which weights it may touch). | model | DiT / UNet | text encoder(s) | VAE | |---|---|---|---| -| FLUX.1 (12B) | nf4, fp8 (streamed), bf16, fp32 | T5: fp32 (CPU) · CLIP-L: fp16/fp32 | 16/32 | +| FLUX.1 (12B) | nf4, fp8 (streamed), bf16, fp32 | T5: bf16 (GPU, 14 GB+) or fp32 (CPU) · CLIP-L: fp16/fp32 | 16/32 | | FLUX.2 klein (4B) | nf4, fp8 (resident), bf16, fp32 | Qwen3-4B: bf16/fp32 | 16/32 | | Z-Image (6B) | nf4, fp8 (resident), bf16, fp32 | Qwen3-4B: bf16/fp32 | 16/32 | | LTX-2.3 (22B video) | nf4, fp8 (streamed), bf16, fp32 | Gemma3-12B: nf4, bf16, fp32 | 16/32 | @@ -57,10 +58,11 @@ device maps (`auto_devices()` strategies: `full_gpu`, `unet_gpu`, `cpu_only`); the flux family and LTX use phase offloading, where each component holds the GPU only for its own phase — text encoding, denoising, decoding — and the denoiser is the sole GPU tenant during -the loop. Text encoders earn special placement: FLUX.1's T5 runs -fp32 on the CPU for quality; the Qwen3 encoders phase-onload in bf16; -the Gemma3 encoder can be GPU-resident or CPU-resident with a staged -swap (see below). +the loop. Text encoders earn special placement: FLUX.1's T5 +phase-onloads in bf16 on 14 GB+ cards (its ~9.8 GB encode phase fits) +and runs fp32 on the CPU below that; the Qwen3 encoders phase-onload in +bf16; the Gemma3 encoder can be GPU-resident or CPU-resident with a +staged swap (see below). ## Lever 3: residency @@ -78,11 +80,12 @@ From most to least VRAM: (11 GB of transformer re-onloads in 0.5 s), and the offload is a pointer swap back to the still-valid pinned copy — zero bytes moved, because inference never mutates weights. Page-locking costs - ~0.6 s/GB once at load. Built today for the LTX pipeline - components and the Gemma3 encoder. + ~0.6 s/GB once at load. Built for the LTX pipeline components, the + Gemma3 encoder, and the FLUX-family transformer, VAE decoder, and + text encoder(s). 4. **Pageable phase swap** — the same movement through ordinary - memory, at roughly 2-16 GB/s depending on tensor layout. What the - flux-family encoders do today. + memory, at roughly 2-16 GB/s depending on tensor layout. The + fallback when pinning is opted out or page-locking fails. 5. **Streamed weights** — the bigger-than-VRAM tier: weights stay in (pinned) host RAM permanently and stream across PCIe during each forward pass, about one byte per parameter per step. LTX fp8 and @@ -107,8 +110,10 @@ soft per component. The global switch is `options(diffuseR.pin_staging = FALSE)` — reach for it under host memory pressure, in containers with hard memory caps, or for single-generation sessions where the one-time page-lock never pays -itself back. Today the LTX pipeline and Gemma3 loaders consume the -decision; the image-model loaders do not stage weights yet. +itself back. The LTX pipeline, the Gemma3 loaders, and the FLUX-family +image loaders consume the decision (their `pin` argument defaults to +it); the SD-family loaders place components statically, so pinning is +inert for them. ## Putting it together @@ -121,5 +126,6 @@ r$note # fork suggestion when fp8 wanted but unreadable ``` Treat the result as the machine's advice: pass its fields to the -loaders and generators (only `flux_memory_profile()` and `serve()`'s -LTX artifact selection consume it automatically today). +loaders and generators. The FLUX-family loaders' `pin` (and FLUX.1's +`text_device`), `flux_memory_profile()`, and `serve()`'s LTX artifact +selection consume it automatically.