diff --git a/.gitignore b/.gitignore
index 464ad16..d9d19d1 100644
--- a/.gitignore
+++ b/.gitignore
@@ -51,5 +51,7 @@ rsconnect/
# Upstream reference checkouts and working notes (never committed)
ref/
+/ref
tasks/
.venv/
+tools/cache/
diff --git a/CLAUDE.md b/CLAUDE.md
index f2683e9..54e8094 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -316,10 +316,37 @@ See cornyverse CLAUDE.md for safetensors package setup (use cornball-ai fork unt
- Handles Blackwell workaround automatically
### Model Support
-- [ ] Add FLUX model support
+- [x] Add FLUX model support (FLUX.1-schnell, see below)
- [ ] Add SD3 model support
- [ ] ControlNet integration
+### FLUX.1-schnell (Complete)
+
+4-step distilled text-to-image, CFG-free. 12B MMDiT + T5-XXL + CLIP-L +
+16-channel VAE, ported from diffusers/transformers (provenance in
+`inst/REFERENCES.md`).
+
+```r
+library(diffuseR)
+download_flux1() # gated HF repo: accept license + set HF_TOKEN; ~34 GB
+ # download, one-time NF4 quantize to a 6.8 GB artifact
+txt2img_flux("An astronaut riding a horse on Mars, photorealistic",
+ seed = 7) # or txt2img("...", model_name = "flux1")
+```
+
+Measured on the RTX 5060 Ti 16 GB (NF4 resident, T5 float32 on CPU):
+1024x1024 in ~2 min wall (peak 8.7 GB alloc / 9.0 GB reserved),
+512x512 in ~1.5 min (peak 8.0 GB). CPU-only works too (~9 min at 256px).
+
+Key components: `flux_transformer` (19 double + 38 single blocks),
+`t5_encoder` + `unigram_tokenizer` (pure R SentencePiece Viterbi),
+`flux_quantize`/`flux_load_transformer` (NF4 ~6.8 GB resident or fp8
+~12 GB streamed; cast census is exactly 494 weights), `flux_load_pipeline`
++ `txt2img_flux` with per-phase GPU offloading. Gotchas that bit once:
+the shipped tokenizer_2 uses Metaspace prepend "always" (a spiece
+conversion gives "never" - different ids for every prompt), and quantized
+residents must follow the compute dtype (bf16 GPU / fp32 CPU).
+
### LTX-2.3 Video Generation (clean-room rewrite in progress)
The original LTX-2.0 port was removed and is being replaced by a ground-up
diff --git a/DESCRIPTION b/DESCRIPTION
index 7d344d6..9201130 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,6 +1,6 @@
Package: diffuseR
Title: Functional Interface to Diffusion Models in R
-Version: 0.1.0
+Version: 0.1.0.1
Authors@R: c(
person("Troy", "Hernandez", email = "troy@cornball.ai", role = c("aut", "cre"),
comment = c(ORCID = "0009-0005-4248-604X")),
diff --git a/NAMESPACE b/NAMESPACE
index 7f9d48e..3f00af7 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -3,31 +3,57 @@
export(auto_devices)
export(bpe_tokenizer)
export(clear_vram)
+export(clip_pooled_output)
export(CLIPTokenizer)
export(ddim_scheduler_create)
export(ddim_scheduler_step)
export(decode_bpe)
export(download_component)
+export(download_flux1)
export(download_ltx2)
export(download_model)
export(encode_bpe)
+export(encode_unigram)
export(encode_with_gemma3)
+export(encode_with_t5)
export(filename_from_prompt)
export(flowmatch_calculate_shift)
export(flowmatch_scale_noise)
export(flowmatch_scheduler_create)
export(flowmatch_scheduler_step)
export(flowmatch_set_timesteps)
+export(flux_ada_layer_norm_continuous)
+export(flux_ada_layer_norm_zero)
+export(flux_ada_layer_norm_zero_single)
+export(flux_apply_rotary_emb)
+export(flux_attention)
+export(flux_double_block)
+export(flux_is_quant_key)
+export(flux_load_pipeline)
+export(flux_load_transformer)
+export(flux_memory_profile)
+export(flux_open_checkpoint)
+export(flux_open_quantized)
+export(flux_pack_latents)
+export(flux_pos_embed)
+export(flux_prepare_latent_image_ids)
+export(flux_quantize)
+export(flux_single_block)
+export(flux_transformer)
+export(flux_unpack_latents)
export(gemma3_config_ltx2)
export(gemma3_text_model)
export(gemma3_tokenizer)
export(img2img)
export(is_blackwell_gpu)
export(latents_to_video)
+export(load_decoder_safetensors)
export(load_decoder_weights)
export(load_gemma3_text_encoder)
export(load_model_component)
export(load_pipeline)
+export(load_t5_text_encoder)
+export(load_text_encoder_safetensors)
export(load_text_encoder_weights)
export(load_text_encoder2_weights)
export(load_to_gpu)
@@ -124,10 +150,12 @@ export(save_video)
export(save_video_ltx23)
export(scheduler_add_noise)
export(sdxl_memory_profile)
+export(t5_encoder)
export(text_encoder_native)
export(text_encoder2_native)
export(tokenize_gemma3)
export(txt2img)
+export(txt2img_flux)
export(txt2img_sd21)
export(txt2img_sdxl)
export(txt2vid_ltx2)
@@ -135,6 +163,7 @@ export(unet_native)
export(unet_native_from_torchscript)
export(unet_sdxl_native)
export(unet_sdxl_native_from_torchscript)
+export(unigram_tokenizer)
export(vae_decoder_native)
export(vocab_size)
export(vram_report)
@@ -142,5 +171,6 @@ export(write_wav)
S3method(print,bpe_tokenizer)
S3method(print,ltx23_checkpoint)
+S3method(print,unigram_tokenizer)
importFrom(utils,head)
diff --git a/R/checkpoint_flux.R b/R/checkpoint_flux.R
new file mode 100644
index 0000000..841df5c
--- /dev/null
+++ b/R/checkpoint_flux.R
@@ -0,0 +1,141 @@
+#' FLUX Checkpoint Readers
+#'
+#' FLUX transformers ship in the diffusers layout: a directory with
+#' \code{config.json}, one or more \code{diffusion_pytorch_model*.safetensors}
+#' shards, and (when sharded) a
+#' \code{diffusion_pytorch_model.safetensors.index.json} weight map.
+#' These helpers open that layout behind the same checkpoint interface as
+#' \code{\link{ltx23_open_checkpoint}}, so the LTX group loaders and
+#' quantization machinery work unchanged. FLUX module names mirror the
+#' checkpoint keys 1:1 - no key mapping is needed.
+#'
+#' @name checkpoint_flux
+NULL
+
+# Quantization cast set: every large linear weight in the FLUX blocks,
+# including the adaLN modulation linears (norm*.linear are 3.2B params
+# across the model; leaving them bf16 would not fit resident on 16 GB).
+# Everything else (embedders, q/k norms, final norm, biases) stays bf16.
+# Full-size census: 19 double blocks x 14 + 38 single blocks x 6 = 494
+# cast weights, ~11.8B of the 12B parameters.
+.flux_quant_cast_pattern <- paste0(
+ "^(",
+ "transformer_blocks\\.[0-9]+\\.(",
+ "attn\\.(to_q|to_k|to_v|add_q_proj|add_k_proj|add_v_proj|to_out\\.0|to_add_out)",
+ "|ff\\.net\\.(0\\.proj|2)|ff_context\\.net\\.(0\\.proj|2)",
+ "|norm1\\.linear|norm1_context\\.linear",
+ ")",
+ "|single_transformer_blocks\\.[0-9]+\\.(",
+ "attn\\.(to_q|to_k|to_v)|proj_mlp|proj_out|norm\\.linear",
+ ")",
+ ")\\.weight$"
+)
+
+#' Test whether a FLUX key is in the quantization cast set
+#'
+#' @param key Character vector of parameter names (diffusers-style).
+#'
+#' @return Logical vector.
+#'
+#' @export
+flux_is_quant_key <- function(key) {
+ grepl(.flux_quant_cast_pattern, key)
+}
+
+#' Open a FLUX transformer checkpoint directory
+#'
+#' Opens a diffusers-layout transformer directory lazily (headers only).
+#' Sharded checkpoints are resolved through the index.json weight map;
+#' single-file checkpoints are opened directly. The transformer
+#' \code{config.json} is attached as \code{$config}.
+#'
+#' @param transformer_dir Directory containing \code{config.json} and the
+#' \code{diffusion_pytorch_model*.safetensors} file(s).
+#'
+#' @return An object of class \code{ltx23_checkpoint} (shared checkpoint
+#' interface): list with \code{handle$get_tensor}, \code{keys},
+#' \code{config}, and \code{path}.
+#'
+#' @export
+flux_open_checkpoint <- function(transformer_dir) {
+ if (!requireNamespace("safetensors", quietly = TRUE)) {
+ stop("The safetensors package is required to read FLUX checkpoints.")
+ }
+ transformer_dir <- path.expand(transformer_dir)
+ if (!dir.exists(transformer_dir)) {
+ stop("Checkpoint directory not found: ", transformer_dir)
+ }
+
+ config <- NULL
+ config_path <- file.path(transformer_dir, "config.json")
+ if (file.exists(config_path)) {
+ config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE)
+ }
+
+ opened <- .flux_open_sharded_dir(transformer_dir, "diffusion_pytorch_model")
+
+ structure(
+ list(handle = opened$handle, keys = opened$keys, version = NULL,
+ config = config, path = transformer_dir),
+ class = "ltx23_checkpoint"
+ )
+}
+
+# Open a HF-layout safetensors directory (sharded via .safetensors
+# .index.json, or a single .safetensors) lazily; returns the
+# handle/keys pair shared by all checkpoint objects.
+.flux_open_sharded_dir <- function(dir, base) {
+ index_path <- file.path(dir, paste0(base, ".safetensors.index.json"))
+ if (file.exists(index_path)) {
+ index <- jsonlite::fromJSON(index_path, simplifyVector = TRUE)
+ weight_map <- unlist(index$weight_map)
+ shard_files <- unique(weight_map)
+ missing <- shard_files[!file.exists(file.path(dir, shard_files))]
+ if (length(missing)) {
+ stop("Missing checkpoint shards: ", paste(missing, collapse = ", "))
+ }
+ handles <- lapply(file.path(dir, shard_files), function(p) {
+ safetensors::safetensors$new(p, framework = "torch")
+ })
+ names(handles) <- shard_files
+ keys <- names(weight_map)
+ handle <- list(
+ get_tensor = function(key) {
+ shard <- weight_map[[key]]
+ if (is.null(shard) || is.na(shard)) {
+ stop("Key not found in checkpoint index: ", key)
+ }
+ handles[[shard]]$get_tensor(key)
+ }
+ )
+ } else {
+ single_path <- file.path(dir, paste0(base, ".safetensors"))
+ if (!file.exists(single_path)) {
+ stop("No ", base, " safetensors (or index) in ", dir)
+ }
+ h <- safetensors::safetensors$new(single_path, framework = "torch")
+ keys <- setdiff(h$keys(), "__metadata__")
+ handle <- list(get_tensor = function(key) h$get_tensor(key))
+ }
+ list(handle = handle, keys = keys)
+}
+
+#' Open a quantized FLUX artifact directory
+#'
+#' Opens the sharded NF4/fp8 artifact written by
+#' \code{\link{flux_quantize}} through the shared checkpoint interface.
+#' The manifest's embedded transformer config and \code{format} ride
+#' along, so \code{\link{flux_load_transformer}} needs nothing else.
+#'
+#' @param dir The quantized artifact directory (with manifest.json).
+#'
+#' @return An \code{ltx23_checkpoint} with \code{$format} set.
+#'
+#' @export
+flux_open_quantized <- function(dir) {
+ manifest_path <- file.path(dir, "manifest.json")
+ if (!file.exists(manifest_path)) {
+ stop("No manifest.json in ", dir, "; run flux_quantize() first.")
+ }
+ ltx23_open_fp8_checkpoint(dir)
+}
diff --git a/R/dit_flux.R b/R/dit_flux.R
new file mode 100644
index 0000000..1c2a28f
--- /dev/null
+++ b/R/dit_flux.R
@@ -0,0 +1,151 @@
+#' FLUX Transformer (MMDiT)
+#'
+#' Fresh R port of FluxTransformer2DModel from the diffusers reference
+#' implementation (Apache-2.0,
+#' src/diffusers/models/transformers/transformer_flux.py). The module
+#' tree mirrors the diffusers state-dict keys 1:1, so checkpoints load
+#' without remapping. FLUX.1-schnell has no guidance embedder
+#' (guidance_embeds = FALSE); the guidance-distilled dev variant is not
+#' implemented.
+#'
+#' @name dit_flux
+NULL
+
+# Combined timestep + pooled-text conditioning, matching diffusers
+# CombinedTimestepTextProjEmbeddings state-dict names. Both embedders are
+# linear_1 -> silu -> linear_2, which ltx23_timestep_embedding provides
+# (PixArtAlphaTextProjection with act_fn = "silu" is the same function).
+flux_time_text_embed <- torch::nn_module(
+ "flux_time_text_embed",
+ initialize = function(embedding_dim, pooled_projection_dim) {
+ self$timestep_embedder <- ltx23_timestep_embedding(256L, embedding_dim)
+ self$text_embedder <- ltx23_timestep_embedding(pooled_projection_dim,
+ embedding_dim)
+},
+ forward = function(timestep, pooled_projection) {
+ proj <- ltx23_get_timestep_embedding(timestep, 256L,
+ flip_sin_to_cos = TRUE, downscale_freq_shift = 0)
+ temb <- self$timestep_embedder(proj$to(dtype = pooled_projection$dtype))
+ temb + self$text_embedder(pooled_projection)
+}
+)
+
+#' FLUX transformer model
+#'
+#' 19 double-stream (MMDiT) blocks followed by 38 single-stream blocks
+#' over the joint [text; image] sequence, with adaLN-Zero conditioning on
+#' timestep + pooled CLIP text. Rotary embeddings are precomputed by the
+#' caller with \code{flux_pos_embed} (they are static across denoise
+#' steps). Defaults are the FLUX.1-schnell configuration.
+#'
+#' @param in_channels Integer. Packed latent channels (64).
+#' @param num_layers Integer. Double-stream block count.
+#' @param num_single_layers Integer. Single-stream block count.
+#' @param attention_head_dim Integer. Per-head dimension.
+#' @param num_attention_heads Integer. Attention heads.
+#' @param joint_attention_dim Integer. T5 embedding dim (4096).
+#' @param pooled_projection_dim Integer. CLIP pooled dim (768).
+#' @param axes_dims_rope Integer vector. Per-axis rotary dims.
+#' @param out_channels Integer or NULL. Output channels (defaults to
+#' \code{in_channels}).
+#'
+#' @return Module whose forward(hidden_states, encoder_hidden_states,
+#' pooled_projections, timestep, image_rotary_emb) returns the
+#' predicted velocity for the image tokens [B, S_img, out_channels].
+#' \code{timestep} is in sigma space (0-1); it is scaled by 1000
+#' internally, matching the reference.
+#'
+#' @export
+flux_transformer <- torch::nn_module(
+ "flux_transformer",
+ initialize = function(in_channels = 64L,
+ num_layers = 19L,
+ num_single_layers = 38L,
+ attention_head_dim = 128L,
+ num_attention_heads = 24L,
+ joint_attention_dim = 4096L,
+ pooled_projection_dim = 768L,
+ axes_dims_rope = c(16L, 56L, 56L),
+ out_channels = NULL) {
+ inner_dim <- num_attention_heads * attention_head_dim
+ self$inner_dim <- inner_dim
+ self$axes_dims_rope <- as.integer(axes_dims_rope)
+ self$out_channels <- as.integer(out_channels %||% in_channels)
+
+ self$time_text_embed <- flux_time_text_embed(inner_dim,
+ pooled_projection_dim)
+ self$context_embedder <- torch::nn_linear(joint_attention_dim, inner_dim)
+ self$x_embedder <- torch::nn_linear(in_channels, inner_dim)
+
+ self$transformer_blocks <- torch::nn_module_list(
+ lapply(seq_len(num_layers), function(i) {
+ flux_double_block(inner_dim, num_attention_heads, attention_head_dim)
+ })
+ )
+ self$single_transformer_blocks <- torch::nn_module_list(
+ lapply(seq_len(num_single_layers), function(i) {
+ flux_single_block(inner_dim, num_attention_heads, attention_head_dim)
+ })
+ )
+
+ self$norm_out <- flux_ada_layer_norm_continuous(inner_dim, inner_dim)
+ self$proj_out <- torch::nn_linear(inner_dim, self$out_channels, bias = TRUE)
+},
+ forward = function(hidden_states, encoder_hidden_states,
+ pooled_projections, timestep, image_rotary_emb,
+ chunk_size = NULL) {
+ hidden_states <- self$x_embedder(hidden_states)
+ timestep <- timestep$to(dtype = hidden_states$dtype)$mul(1000)
+ temb <- self$time_text_embed(timestep, pooled_projections)
+ encoder_hidden_states <- self$context_embedder(encoder_hidden_states)
+
+ block_gc <- isTRUE(getOption("diffuseR.block_gc"))
+ debug <- isTRUE(getOption("diffuseR.debug"))
+
+ for (i in seq_along(self$transformer_blocks)) {
+ res <- self$transformer_blocks[[i]](
+ hidden_states = hidden_states,
+ encoder_hidden_states = encoder_hidden_states,
+ temb = temb,
+ image_rotary_emb = image_rotary_emb,
+ chunk_size = chunk_size
+ )
+ encoder_hidden_states <- res[[1]]
+ hidden_states <- res[[2]]
+ if (debug && torch::cuda_is_available()) {
+ ms <- torch::cuda_memory_stats()
+ message(sprintf(" double block %d: %.2f GB allocated", i,
+ ms$allocated_bytes$all$current / 1e9))
+ }
+ if (block_gc) {
+ gc(verbose = FALSE)
+ }
+ }
+
+ # The reference concatenates [text; image] inside every single block
+ # and splits after; concatenating once here is numerically identical
+ txt_len <- encoder_hidden_states$shape[2]
+ hidden_states <- torch::torch_cat(
+ list(encoder_hidden_states, hidden_states),
+ dim = 2L
+ )
+ for (i in seq_along(self$single_transformer_blocks)) {
+ hidden_states <- self$single_transformer_blocks[[i]](
+ hidden_states = hidden_states,
+ temb = temb,
+ image_rotary_emb = image_rotary_emb,
+ chunk_size = chunk_size
+ )
+ if (block_gc) {
+ gc(verbose = FALSE)
+ }
+ }
+ hidden_states <- hidden_states$narrow(
+ 2L, txt_len + 1L,
+ hidden_states$shape[2] - txt_len
+ )
+
+ hidden_states <- self$norm_out(hidden_states, temb)
+ self$proj_out(hidden_states)
+}
+)
diff --git a/R/dit_flux_modules.R b/R/dit_flux_modules.R
new file mode 100644
index 0000000..959a17e
--- /dev/null
+++ b/R/dit_flux_modules.R
@@ -0,0 +1,307 @@
+#' FLUX Transformer Building Blocks
+#'
+#' Fresh R port of the FLUX MMDiT blocks from the diffusers reference
+#' implementation (Apache-2.0,
+#' src/diffusers/models/transformers/transformer_flux.py and
+#' src/diffusers/models/normalization.py). Module field names mirror the
+#' diffusers state-dict keys 1:1 so checkpoints load without remapping.
+#' Reuses the LTX primitives \code{ltx23_rms_norm}, \code{.ltx23_sdpa}
+#' and \code{ltx23_feed_forward}.
+#'
+#' @name dit_flux_modules
+NULL
+
+#' FLUX adaLN-Zero modulation (double-stream)
+#'
+#' Projects the conditioning embedding to six modulation vectors and
+#' returns the msa-modulated input plus the remaining parameters.
+#' Reference: diffusers AdaLayerNormZero.
+#'
+#' @param dim Integer. Model dimension.
+#'
+#' @return Module whose forward(x, emb) returns
+#' \code{list(x_norm, gate_msa, shift_mlp, scale_mlp, gate_mlp)}.
+#'
+#' @export
+flux_ada_layer_norm_zero <- torch::nn_module(
+ "flux_ada_layer_norm_zero",
+ initialize = function(dim) {
+ self$linear <- torch::nn_linear(dim, 6L * dim, bias = TRUE)
+ self$norm <- torch::nn_layer_norm(dim, eps = 1e-6,
+ elementwise_affine = FALSE)
+},
+ forward = function(x, emb) {
+ emb <- self$linear(torch::nnf_silu(emb))
+ p <- emb$chunk(6L, dim = 2L)
+ # shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp
+ x <- self$norm(x) * p[[2]]$unsqueeze(2L)$add(1) + p[[1]]$unsqueeze(2L)
+ list(x, p[[3]], p[[4]], p[[5]], p[[6]])
+}
+)
+
+#' FLUX adaLN-Zero modulation (single-stream)
+#'
+#' Three modulation vectors: shift, scale, gate. Reference: diffusers
+#' AdaLayerNormZeroSingle.
+#'
+#' @param dim Integer. Model dimension.
+#'
+#' @return Module whose forward(x, emb) returns \code{list(x_norm, gate)}.
+#'
+#' @export
+flux_ada_layer_norm_zero_single <- torch::nn_module(
+ "flux_ada_layer_norm_zero_single",
+ initialize = function(dim) {
+ self$linear <- torch::nn_linear(dim, 3L * dim, bias = TRUE)
+ self$norm <- torch::nn_layer_norm(dim, eps = 1e-6,
+ elementwise_affine = FALSE)
+},
+ forward = function(x, emb) {
+ emb <- self$linear(torch::nnf_silu(emb))
+ p <- emb$chunk(3L, dim = 2L)
+ # shift_msa, scale_msa, gate_msa
+ x <- self$norm(x) * p[[2]]$unsqueeze(2L)$add(1) + p[[1]]$unsqueeze(2L)
+ list(x, p[[3]])
+}
+)
+
+#' FLUX continuous adaLN (final norm)
+#'
+#' Scale/shift conditioning of the final norm. Note the chunk order:
+#' scale first, then shift (the reverse of adaLN-Zero). Reference:
+#' diffusers AdaLayerNormContinuous as used by FLUX norm_out
+#' (elementwise_affine = FALSE, eps = 1e-6).
+#'
+#' @param dim Integer. Model dimension.
+#' @param cond_dim Integer. Conditioning embedding dimension.
+#'
+#' @export
+flux_ada_layer_norm_continuous <- torch::nn_module(
+ "flux_ada_layer_norm_continuous",
+ initialize = function(dim, cond_dim = dim) {
+ self$linear <- torch::nn_linear(cond_dim, 2L * dim, bias = TRUE)
+ self$norm <- torch::nn_layer_norm(dim, eps = 1e-6,
+ elementwise_affine = FALSE)
+},
+ forward = function(x, cond) {
+ emb <- self$linear(torch::nnf_silu(cond)$to(dtype = x$dtype))
+ p <- emb$chunk(2L, dim = 2L)
+ # scale, shift
+ self$norm(x) * p[[1]]$unsqueeze(2L)$add(1) + p[[2]]$unsqueeze(2L)
+}
+)
+
+#' FLUX joint attention
+#'
+#' Multi-head attention with per-head RMS q/k norms and rotary position
+#' embeddings. With \code{added_kv = TRUE} (double-stream blocks) the
+#' text stream gets its own q/k/v projections and both streams attend
+#' jointly (text tokens first); the outputs are split back and projected
+#' per stream. With \code{pre_only = TRUE} (single-stream blocks) there
+#' is no output projection. Reference: diffusers FluxAttention +
+#' FluxAttnProcessor.
+#'
+#' @param query_dim Integer. Model dimension.
+#' @param heads Integer. Number of attention heads.
+#' @param dim_head Integer. Per-head dimension.
+#' @param added_kv Logical. Add text-stream projections (double blocks).
+#' @param pre_only Logical. Skip the output projection (single blocks).
+#' @param eps Numeric. RMS norm epsilon.
+#'
+#' @export
+flux_attention <- torch::nn_module(
+ "flux_attention",
+ initialize = function(query_dim, heads, dim_head, added_kv = FALSE,
+ pre_only = FALSE, eps = 1e-6) {
+ inner_dim <- heads * dim_head
+ self$heads <- heads
+ self$dim_head <- dim_head
+ self$added_kv <- added_kv
+ self$pre_only <- pre_only
+
+ self$norm_q <- ltx23_rms_norm(dim_head, eps = eps)
+ self$norm_k <- ltx23_rms_norm(dim_head, eps = eps)
+ self$to_q <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+ self$to_k <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+ self$to_v <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+
+ if (!pre_only) {
+ self$to_out <- torch::nn_module_list(list(
+ torch::nn_linear(inner_dim, query_dim, bias = TRUE)
+ ))
+ }
+ if (added_kv) {
+ self$norm_added_q <- ltx23_rms_norm(dim_head, eps = eps)
+ self$norm_added_k <- ltx23_rms_norm(dim_head, eps = eps)
+ self$add_q_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+ self$add_k_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+ self$add_v_proj <- torch::nn_linear(query_dim, inner_dim, bias = TRUE)
+ self$to_add_out <- torch::nn_linear(inner_dim, query_dim, bias = TRUE)
+ }
+},
+ forward = function(hidden_states, encoder_hidden_states = NULL,
+ image_rotary_emb = NULL, chunk_size = NULL) {
+ # Per-head layout [B, S, H, D]
+ query <- self$to_q(hidden_states)$unflatten(3L, c(self$heads, -1L))
+ key <- self$to_k(hidden_states)$unflatten(3L, c(self$heads, -1L))
+ value <- self$to_v(hidden_states)$unflatten(3L, c(self$heads, -1L))
+
+ query <- self$norm_q(query)
+ key <- self$norm_k(key)
+
+ if (!is.null(encoder_hidden_states)) {
+ txt_len <- encoder_hidden_states$shape[2]
+ eq <- self$add_q_proj(encoder_hidden_states)$unflatten(3L,
+ c(self$heads, -1L))
+ ek <- self$add_k_proj(encoder_hidden_states)$unflatten(3L, c(self$heads, -1L))
+ ev <- self$add_v_proj(encoder_hidden_states)$unflatten(3L, c(self$heads, -1L))
+ eq <- self$norm_added_q(eq)
+ ek <- self$norm_added_k(ek)
+ # Text tokens first, matching the rotary frequency layout
+ query <- torch::torch_cat(list(eq, query), dim = 2L)
+ key <- torch::torch_cat(list(ek, key), dim = 2L)
+ value <- torch::torch_cat(list(ev, value), dim = 2L)
+ }
+
+ # [B, S, H, D] -> [B, H, S, D] for RoPE and attention
+ query <- query$transpose(2L, 3L)
+ key <- key$transpose(2L, 3L)
+ value <- value$transpose(2L, 3L)
+
+ if (!is.null(image_rotary_emb)) {
+ query <- flux_apply_rotary_emb(query, image_rotary_emb)
+ key <- flux_apply_rotary_emb(key, image_rotary_emb)
+ }
+
+ out <- .ltx23_sdpa(query, key, value, chunk_size = chunk_size)
+ # [B, H, S, D] -> [B, S, H*D]
+ out <- out$transpose(2L, 3L)$flatten(start_dim = 3L)
+ out <- out$to(dtype = hidden_states$dtype)
+
+ if (!is.null(encoder_hidden_states)) {
+ seq_len <- out$shape[2]
+ ctx <- out$narrow(2L, 1L, txt_len)
+ img <- out$narrow(2L, txt_len + 1L, seq_len - txt_len)
+ return(list(
+ self$to_out[[1]](img$contiguous()),
+ self$to_add_out(ctx$contiguous())
+ ))
+ }
+ if (self$pre_only) {
+ return(out)
+ }
+ self$to_out[[1]](out)
+}
+)
+
+#' FLUX double-stream (MMDiT) transformer block
+#'
+#' Image and text streams each get adaLN-Zero modulation and a
+#' feed-forward; attention is joint across both streams. Reference:
+#' diffusers FluxTransformerBlock.
+#'
+#' @param dim Integer. Model dimension.
+#' @param num_attention_heads Integer. Attention heads.
+#' @param attention_head_dim Integer. Per-head dimension.
+#'
+#' @return Module whose forward(hidden_states, encoder_hidden_states,
+#' temb, image_rotary_emb) returns
+#' \code{list(encoder_hidden_states, hidden_states)}.
+#'
+#' @export
+flux_double_block <- torch::nn_module(
+ "flux_double_block",
+ initialize = function(dim, num_attention_heads, attention_head_dim) {
+ self$norm1 <- flux_ada_layer_norm_zero(dim)
+ self$norm1_context <- flux_ada_layer_norm_zero(dim)
+ self$attn <- flux_attention(dim, num_attention_heads, attention_head_dim,
+ added_kv = TRUE)
+ self$norm2 <- torch::nn_layer_norm(dim, eps = 1e-6,
+ elementwise_affine = FALSE)
+ self$ff <- ltx23_feed_forward(dim, mult = 4L)
+ self$norm2_context <- torch::nn_layer_norm(dim, eps = 1e-6,
+ elementwise_affine = FALSE)
+ self$ff_context <- ltx23_feed_forward(dim, mult = 4L)
+},
+ forward = function(hidden_states, encoder_hidden_states, temb,
+ image_rotary_emb = NULL, chunk_size = NULL) {
+ n1 <- self$norm1(hidden_states, emb = temb)
+ n1c <- self$norm1_context(encoder_hidden_states, emb = temb)
+
+ attn_out <- self$attn(
+ hidden_states = n1[[1]],
+ encoder_hidden_states = n1c[[1]],
+ image_rotary_emb = image_rotary_emb,
+ chunk_size = chunk_size
+ )
+
+ # Image stream: gated attention + modulated feed-forward
+ hidden_states <- hidden_states + n1[[2]]$unsqueeze(2L) * attn_out[[1]]
+ norm_h <- self$norm2(hidden_states) * n1[[4]]$unsqueeze(2L)$add(1) +
+ n1[[3]]$unsqueeze(2L)
+ hidden_states <- hidden_states + n1[[5]]$unsqueeze(2L) * self$ff(norm_h)
+
+ # Text stream mirrors with its own modulation
+ encoder_hidden_states <- encoder_hidden_states +
+ n1c[[2]]$unsqueeze(2L) * attn_out[[2]]
+ norm_c <- self$norm2_context(encoder_hidden_states) *
+ n1c[[4]]$unsqueeze(2L)$add(1) + n1c[[3]]$unsqueeze(2L)
+ encoder_hidden_states <- encoder_hidden_states +
+ n1c[[5]]$unsqueeze(2L) * self$ff_context(norm_c)
+
+ if (encoder_hidden_states$dtype == torch::torch_float16()) {
+ encoder_hidden_states <- encoder_hidden_states$clamp(-65504, 65504)
+ }
+ list(encoder_hidden_states, hidden_states)
+}
+)
+
+#' FLUX single-stream transformer block
+#'
+#' Parallel attention + MLP over the joint [text; image] sequence with a
+#' shared gate: \code{x + gate * proj_out(cat(attn, gelu(mlp)))}. The
+#' reference concatenates the streams inside every block and splits after;
+#' here the caller concatenates once before the single-block stack, which
+#' is numerically identical. Reference: diffusers
+#' FluxSingleTransformerBlock.
+#'
+#' @param dim Integer. Model dimension.
+#' @param num_attention_heads Integer. Attention heads.
+#' @param attention_head_dim Integer. Per-head dimension.
+#' @param mlp_ratio Numeric. MLP hidden dim multiplier.
+#'
+#' @return Module whose forward(hidden_states, temb, image_rotary_emb)
+#' returns the joint hidden states.
+#'
+#' @export
+flux_single_block <- torch::nn_module(
+ "flux_single_block",
+ initialize = function(dim, num_attention_heads, attention_head_dim,
+ mlp_ratio = 4.0) {
+ mlp_hidden_dim <- as.integer(dim * mlp_ratio)
+ self$norm <- flux_ada_layer_norm_zero_single(dim)
+ self$proj_mlp <- torch::nn_linear(dim, mlp_hidden_dim)
+ self$proj_out <- torch::nn_linear(dim + mlp_hidden_dim, dim)
+ self$attn <- flux_attention(dim, num_attention_heads, attention_head_dim,
+ pre_only = TRUE)
+},
+ forward = function(hidden_states, temb, image_rotary_emb = NULL,
+ chunk_size = NULL) {
+ residual <- hidden_states
+ n <- self$norm(hidden_states, emb = temb)
+ mlp <- torch::nnf_gelu(self$proj_mlp(n[[1]]), approximate = "tanh")
+ attn_out <- self$attn(
+ hidden_states = n[[1]],
+ image_rotary_emb = image_rotary_emb,
+ chunk_size = chunk_size
+ )
+
+ # Attention half first, then the MLP half
+ hidden_states <- torch::torch_cat(list(attn_out, mlp), dim = 3L)
+ hidden_states <- residual + n[[2]]$unsqueeze(2L) * self$proj_out(hidden_states)
+ if (hidden_states$dtype == torch::torch_float16()) {
+ hidden_states <- hidden_states$clamp(-65504, 65504)
+ }
+ hidden_states
+}
+)
diff --git a/R/download_flux.R b/R/download_flux.R
new file mode 100644
index 0000000..9c22c27
--- /dev/null
+++ b/R/download_flux.R
@@ -0,0 +1,168 @@
+#' Download and Prepare FLUX.1-schnell Weights
+#'
+#' Downloads FLUX.1-schnell from HuggingFace (weights Apache-2.0, but
+#' the repo is gated behind a license click-through) and quantizes the
+#' 12B transformer to a local NF4 (~7 GB) or fp8 (~12 GB) artifact.
+#'
+#' @name download_flux
+NULL
+
+.flux1_repo <- "black-forest-labs/FLUX.1-schnell"
+
+.flux1_transformer_files <- c(
+ "transformer/config.json",
+ "transformer/diffusion_pytorch_model.safetensors.index.json",
+ sprintf("transformer/diffusion_pytorch_model-%05d-of-00003.safetensors",
+ 1:3)
+)
+
+.flux1_support_files <- c("vae/config.json",
+ "vae/diffusion_pytorch_model.safetensors",
+ "text_encoder/config.json",
+ "text_encoder/model.safetensors",
+ "scheduler/scheduler_config.json")
+
+.flux1_t5_files <- c(
+ "text_encoder_2/config.json",
+ "text_encoder_2/model.safetensors.index.json",
+ sprintf("text_encoder_2/model-%05d-of-00002.safetensors", 1:2),
+ "tokenizer_2/tokenizer.json",
+ "tokenizer_2/tokenizer_config.json",
+ "tokenizer_2/special_tokens_map.json"
+)
+
+# hub_download with the gated-repo 401/403 turned into an actionable error
+.flux1_download <- function(file, ...) {
+ tryCatch(
+ hfhub::hub_download(.flux1_repo, file, ...),
+ error = function(e) {
+ msg <- conditionMessage(e)
+ if (grepl("401|403|[Uu]nauthorized|[Ff]orbidden", msg)) {
+ stop(
+ "FLUX.1-schnell is a gated HuggingFace repo (the weights are ",
+ "Apache-2.0; the gate is a license click-through). To download:\n",
+ " 1. Log in at https://huggingface.co/black-forest-labs/FLUX.1-schnell ",
+ "and accept the license.\n",
+ " 2. Create a read token at https://huggingface.co/settings/tokens\n",
+ " 3. Sys.setenv(HF_TOKEN = \"hf_...\") and retry.",
+ call. = FALSE
+ )
+ }
+ stop(e)
+ }
+ )
+}
+
+#' Download FLUX.1-schnell and build the quantized artifact
+#'
+#' Skips work already done: a valid quantized manifest short-circuits
+#' the transformer download; cached files are not re-fetched. Needs
+#' \code{HF_TOKEN} set for the gated repo (see the error message it
+#' raises without one). The bf16 transformer source (~24 GB in the
+#' HuggingFace cache) may be deleted after quantization.
+#'
+#' @param quantize Logical. Build the quantized artifact after
+#' downloading.
+#' @param precision "nf4" (~7 GB, GPU-resident on 16 GB cards) or
+#' "fp8" (~12 GB, CPU-resident, streamed; near-bf16 quality).
+#' @param output_dir Directory for the quantized artifact.
+#' @param text_encoders Logical. Also fetch the CLIP + T5 text encoders,
+#' tokenizer, VAE, and scheduler config (~10 GB).
+#' @param verbose Logical.
+#'
+#' @return Invisibly, a list with \code{transformer_dir},
+#' \code{artifact_dir}, and \code{support} (named file paths).
+#'
+#' @export
+download_flux1 <- function(quantize = TRUE, precision = c("nf4", "fp8"),
+ output_dir = NULL, text_encoders = TRUE,
+ verbose = TRUE) {
+ precision <- match.arg(precision)
+ if (is.null(output_dir)) {
+ output_dir <- file.path(tools::R_user_dir("diffuseR", "data"),
+ paste0("flux1-schnell-", precision))
+ }
+ if (!requireNamespace("hfhub", quietly = TRUE)) {
+ stop("The hfhub package is required to download model weights.")
+ }
+ result <- list(transformer_dir = NULL, artifact_dir = output_dir,
+ support = character(0))
+
+ manifest_path <- file.path(output_dir, "manifest.json")
+ have_artifact <- file.exists(manifest_path) && {
+ m <- jsonlite::fromJSON(manifest_path)
+ all(file.exists(file.path(output_dir, m$shards)))
+ }
+
+ if (!have_artifact || !quantize) {
+ cached <- tryCatch(
+ hfhub::hub_download(.flux1_repo, .flux1_transformer_files[[3]],
+ local_files_only = TRUE),
+ error = function(e) NULL
+ )
+ if (is.null(cached) && !have_artifact) {
+ free <- .ltx23_disk_free_gb(path.expand("~"))
+ if (!is.na(free) && free < 45) {
+ warning(sprintf(
+ "Only %.0f GB free; the download + %s artifact need ~45 GB.",
+ free, precision
+ ))
+ }
+ ok <- .ltx23_consent(paste0(
+ "FLUX.1-schnell: the 24 GB bf16 transformer plus a ~",
+ if (precision == "nf4") "7" else "12",
+ " GB local ", precision,
+ " artifact (weights Apache-2.0, gated HuggingFace repo)"
+ ))
+ if (!ok) {
+ stop("Download cancelled.", call. = FALSE)
+ }
+ if (verbose) {
+ message("Downloading the FLUX.1-schnell transformer (24 GB)...")
+ }
+ }
+ paths <- vapply(.flux1_transformer_files, .flux1_download, character(1))
+ result$transformer_dir <- dirname(paths[[1]])
+
+ if (quantize && !have_artifact) {
+ if (verbose) {
+ message("Quantizing transformer linears to ", precision,
+ " (one-time)...")
+ }
+ flux_quantize(result$transformer_dir, output_dir,
+ format = precision, verbose = verbose)
+ if (verbose) {
+ message(
+ toupper(precision), " artifact ready: ", output_dir, "\n",
+ "The 24 GB source in the HuggingFace cache may be deleted ",
+ "if you do not need bf16 weights."
+ )
+ }
+ }
+ } else if (verbose) {
+ message(toupper(precision), " artifact already present: ", output_dir)
+ }
+
+ if (text_encoders) {
+ files <- c(.flux1_support_files, .flux1_t5_files)
+ have_t5 <- !is.null(tryCatch(
+ hfhub::hub_download(.flux1_repo, .flux1_t5_files[[3]],
+ local_files_only = TRUE),
+ error = function(e) NULL
+ ))
+ if (!have_t5) {
+ ok <- .ltx23_consent(
+ "the FLUX text encoders, VAE, and tokenizer (~10 GB)"
+ )
+ if (!ok) {
+ stop("Download cancelled.", call. = FALSE)
+ }
+ if (verbose) {
+ message("Downloading text encoders + VAE...")
+ }
+ }
+ result$support <- vapply(files, .flux1_download, character(1))
+ }
+
+ invisible(result)
+}
diff --git a/R/memory_flux.R b/R/memory_flux.R
new file mode 100644
index 0000000..14f8b97
--- /dev/null
+++ b/R/memory_flux.R
@@ -0,0 +1,46 @@
+#' FLUX Memory Profiles
+#'
+#' VRAM-based execution profiles for the FLUX.1-schnell pipeline,
+#' following the LTX-2.3 profile pattern. The 12B transformer runs NF4
+#' (~7 GB, GPU-resident) or fp8 (~12 GB, CPU-resident and streamed);
+#' the T5-XXL text encoder runs float32 on the CPU by default.
+#'
+#' @name memory_flux
+NULL
+
+#' Resolve a FLUX memory profile
+#'
+#' @param vram_gb Numeric or NULL. Available VRAM; auto-detected when
+#' NULL (via gpu.ctl or nvidia-smi).
+#'
+#' @return List with \code{name}, \code{precision} ("nf4"/"fp8"),
+#' \code{attn_chunk}, \code{text_device}, \code{phase_offload}, and
+#' \code{max_pixels} (largest validated image area).
+#'
+#' @export
+flux_memory_profile <- function(vram_gb = NULL) {
+ if (is.null(vram_gb)) {
+ vram_gb <- .detect_vram(use_free = TRUE)
+ if (is.null(vram_gb) || is.na(vram_gb) || vram_gb <= 0) {
+ vram_gb <- 0
+ }
+ }
+
+ if (vram_gb >= 12) {
+ list(name = "high", precision = "nf4", attn_chunk = NULL,
+ text_device = "cpu", phase_offload = TRUE,
+ max_pixels = 1536L * 1536L)
+ } else if (vram_gb >= 9) {
+ list(name = "medium", precision = "nf4", attn_chunk = 2048L,
+ text_device = "cpu", phase_offload = TRUE,
+ max_pixels = 1024L * 1024L)
+ } else if (vram_gb >= 7) {
+ list(name = "low", precision = "fp8", attn_chunk = 1024L,
+ text_device = "cpu", phase_offload = TRUE,
+ max_pixels = 768L * 768L)
+ } else {
+ list(name = "cpu_only", precision = "nf4", attn_chunk = NULL,
+ text_device = "cpu", phase_offload = FALSE,
+ max_pixels = 512L * 512L)
+ }
+}
diff --git a/R/models2devices.R b/R/models2devices.R
index 5af8188..c0e8638 100644
--- a/R/models2devices.R
+++ b/R/models2devices.R
@@ -57,7 +57,8 @@ get_required_components <- function(model_name) {
# "sd15" = c("unet", "decoder", "text_encoder", "encoder"),
"sd21" = c("unet", "decoder", "text_encoder", "encoder"),
"sdxl" = c("unet", "decoder", "text_encoder", "text_encoder2",
- "encoder")
+ "encoder"),
+ "flux1" = c("transformer", "decoder", "text_encoder", "text_encoder2")
# "sd3" = c("transformer", "decoder", "text_encoder", "text_encoder2", "text_encoder3", "encoder"),
# "cascade" = c("prior", "decoder", "text_encoder", "vqgan")
)
diff --git a/R/quantize_flux.R b/R/quantize_flux.R
new file mode 100644
index 0000000..9021d6d
--- /dev/null
+++ b/R/quantize_flux.R
@@ -0,0 +1,329 @@
+#' FLUX Transformer Quantization and Loading
+#'
+#' Quantize the 12B FLUX transformer to NF4 (~7 GB, GPU-resident on
+#' 16 GB cards) or fp8 (~12 GB, CPU-resident and streamed per forward),
+#' and load any format back into \code{\link{flux_transformer}}. Reuses
+#' the LTX-2.3 quantization machinery (\code{ltx23_nf4_quantize},
+#' \code{ltx23_nf4_linear}, \code{ltx23_fp8_linear}); only the cast set
+#' and the diffusers directory layout are FLUX-specific.
+#'
+#' @name quantize_flux
+NULL
+
+.flux_dtype <- function(dtype) {
+ switch(dtype, bfloat16 = torch::torch_bfloat16(),
+ float16 = torch::torch_float16(), float32 = torch::torch_float32(),
+ stop("Unsupported dtype: ", dtype))
+}
+
+# Transformer constructor arguments from a diffusers config.json
+.flux_transformer_args <- function(config) {
+ if (is.null(config)) {
+ return(list())
+ }
+ if (isTRUE(config$guidance_embeds)) {
+ stop("This checkpoint uses guidance embeddings (FLUX.1-dev); ",
+ "only FLUX.1-schnell (guidance_embeds = false) is supported.")
+ }
+ args <- list(
+ in_channels = config$in_channels,
+ num_layers = config$num_layers,
+ num_single_layers = config$num_single_layers,
+ attention_head_dim = config$attention_head_dim,
+ num_attention_heads = config$num_attention_heads,
+ joint_attention_dim = config$joint_attention_dim,
+ pooled_projection_dim = config$pooled_projection_dim,
+ axes_dims_rope = config$axes_dims_rope,
+ out_channels = config$out_channels
+ )
+ # JSON roundtrips turn null into empty lists; drop both
+ args <- Filter(function(x) !is.null(x) && length(x) > 0L, args)
+ lapply(args, function(x) if (is.numeric(x)) as.integer(x) else x)
+}
+
+#' Quantize a FLUX transformer to NF4 or fp8 shards
+#'
+#' Streams the bf16 diffusers checkpoint tensor by tensor. Cast-set
+#' weights (see \code{\link{flux_is_quant_key}}) are stored as NF4
+#' (packed uint8 + \code{_absmax} float32 blocks) or as
+#' float8_e4m3fn with an absmax/448 per-tensor \code{_scale};
+#' everything else is copied through unchanged. The manifest embeds the
+#' transformer config, so the source checkpoint is not needed again
+#' after quantization.
+#'
+#' @param transformer_dir Source diffusers transformer directory.
+#' @param output_dir Output directory for shards + manifest (default:
+#' the per-format location under \code{tools::R_user_dir}).
+#' @param format "nf4" or "fp8".
+#' @param shard_bytes Numeric. Approximate shard size.
+#' @param force Logical. Re-quantize even if a valid manifest exists.
+#' @param verbose Logical.
+#'
+#' @return Invisibly, the manifest list.
+#'
+#' @export
+flux_quantize <- function(transformer_dir, output_dir = NULL,
+ format = c("nf4", "fp8"), shard_bytes = 4e9,
+ force = FALSE, verbose = TRUE) {
+ format <- match.arg(format)
+ if (is.null(output_dir)) {
+ output_dir <- file.path(tools::R_user_dir("diffuseR", "data"),
+ paste0("flux1-schnell-", format))
+ }
+
+ manifest_path <- file.path(output_dir, "manifest.json")
+ if (!force && file.exists(manifest_path)) {
+ manifest <- jsonlite::fromJSON(manifest_path)
+ if (identical(manifest$format, format) &&
+ all(file.exists(file.path(output_dir, manifest$shards)))) {
+ if (verbose) {
+ message(toupper(format), " artifact already present: ",
+ output_dir)
+ }
+ return(invisible(manifest))
+ }
+ }
+ dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)
+
+ ckpt <- flux_open_checkpoint(transformer_dir)
+ if (format == "fp8") {
+ fp8 <- torch::torch_float8_e4m3fn()
+ }
+
+ shard <- list()
+ shard_size <- 0
+ shard_files <- character(0)
+ n_cast <- 0L
+
+ flush_shard <- function() {
+ if (!length(shard)) {
+ return()
+ }
+ fname <- sprintf("flux1-%s-%05d.safetensors", format,
+ length(shard_files) + 1L)
+ safetensors::safe_save_file(shard, file.path(output_dir, fname))
+ shard_files[[length(shard_files) + 1L]] <<- fname
+ if (verbose) {
+ message(sprintf(" wrote %s (%.2f GB, %d tensors)", fname,
+ shard_size / 1e9, length(shard)))
+ }
+ shard <<- list()
+ shard_size <<- 0
+ gc(verbose = FALSE)
+ }
+
+ keys <- ckpt$keys
+ for (i in seq_along(keys)) {
+ key <- keys[[i]]
+ tensor <- ckpt$handle$get_tensor(key)
+
+ if (flux_is_quant_key(key)) {
+ torch::with_no_grad({
+ if (format == "nf4") {
+ q <- ltx23_nf4_quantize(tensor)
+ shard[[key]] <- q$packed
+ shard[[paste0(key, "_absmax")]] <- q$absmax
+ shard_size <- shard_size + prod(tensor$shape) * 0.5625
+ } else {
+ scale <- tensor$abs()$max()$to(dtype = torch::torch_float32())$
+ clamp(min = 1e-12) / 448
+ shard[[key]] <- (tensor$to(dtype = torch::torch_float32()) / scale)$to(dtype = fp8)
+ shard[[paste0(key, "_scale")]] <- scale
+ shard_size <- shard_size + prod(tensor$shape)
+ }
+ })
+ n_cast <- n_cast + 1L
+ } else {
+ shard[[key]] <- tensor
+ shard_size <- shard_size + prod(tensor$shape) * 2
+ }
+ rm(tensor)
+
+ if (shard_size >= shard_bytes) {
+ flush_shard()
+ }
+ if (i %% 200L == 0L) {
+ gc(verbose = FALSE)
+ if (verbose) {
+ message(sprintf(" quantizing %d/%d tensors", i, length(keys)))
+ }
+ }
+ }
+ flush_shard()
+
+ manifest <- list(
+ source = basename(transformer_dir),
+ format = format,
+ shards = shard_files,
+ tensors = length(keys),
+ cast = n_cast,
+ config = ckpt$config
+ )
+ jsonlite::write_json(manifest, manifest_path, auto_unbox = TRUE,
+ pretty = TRUE)
+ if (verbose) {
+ message(sprintf("Quantized %d/%d tensors to %s across %d shards: %s",
+ n_cast, length(keys), format, length(shard_files),
+ output_dir))
+ }
+ invisible(manifest)
+}
+
+#' Load a FLUX transformer from any checkpoint format
+#'
+#' Builds \code{\link{flux_transformer}} from the checkpoint's embedded
+#' config and loads the weights. Dispatches on the checkpoint format:
+#'
+#' \itemize{
+#' \item full precision (\code{\link{flux_open_checkpoint}}): weights
+#' stream into the model in \code{dtype} on \code{device}.
+#' \item \code{"nf4"} (\code{\link{flux_open_quantized}}): cast-set
+#' linears become \code{ltx23_nf4_linear}; the whole model (packed
+#' weights included) moves to \code{device} and stays resident.
+#' \item \code{"fp8"}: cast-set linears become
+#' \code{ltx23_fp8_linear}; fp8 weights stay CPU-resident (optionally
+#' pinned) and stream to \code{device} inside each forward.
+#' }
+#'
+#' @param ckpt A checkpoint from \code{\link{flux_open_checkpoint}} or
+#' \code{\link{flux_open_quantized}}.
+#' @param device Character. Compute device.
+#' @param dtype Character. Model dtype ("bfloat16" or "float32"). For
+#' quantized formats this sets the resident (non-quantized) tensors
+#' and must match the compute dtype: bfloat16 for GPU compute,
+#' float32 for CPU compute.
+#' @param pin Logical. Pin fp8 host memory for faster transfers.
+#' @param verbose Logical.
+#' @param ... Overrides for \code{\link{flux_transformer}} arguments
+#' (tiny test configs).
+#'
+#' @return The loaded \code{flux_transformer} in eval mode.
+#'
+#' @export
+flux_load_transformer <- function(ckpt, device = "cuda", dtype = "bfloat16",
+ pin = TRUE, verbose = TRUE, ...) {
+ stopifnot(inherits(ckpt, "ltx23_checkpoint"))
+ format <- ckpt$format %||% "full"
+
+ args <- utils::modifyList(.flux_transformer_args(ckpt$config), list(...))
+ model <- do.call(flux_transformer, args)
+
+ if (format == "full") {
+ model$to(dtype = .flux_dtype(dtype))
+ res <- ltx23_load_group(ckpt, ckpt$keys, model, verbose = verbose)
+ if (length(res$unmapped) || length(res$unfilled)) {
+ stop("FLUX transformer load: ", length(res$unmapped),
+ " unmapped keys, ", length(res$unfilled), " unfilled params")
+ }
+ model$to(device = device)
+ model$eval()
+ return(model)
+ }
+
+ if (!format %in% c("nf4", "fp8")) {
+ stop("Unknown checkpoint format: ", format)
+ }
+ # Residents (embedders, norms, biases) in the compute dtype: the
+ # quantized linears dequantize into the input's dtype at forward, so
+ # the two must agree (bfloat16 on GPU, float32 for CPU compute)
+ model$to(dtype = .flux_dtype(dtype))
+
+ if (format == "nf4") {
+ sib_suffix <- "_absmax"
+ } else {
+ sib_suffix <- "_scale"
+ }
+ sib_keys <- ckpt$keys[endsWith(ckpt$keys, paste0(".weight", sib_suffix))]
+ main_keys <- setdiff(ckpt$keys, sib_keys)
+
+ dests <- c(model$named_parameters(), model$named_buffers())
+ filled <- character(0)
+ unmapped <- character(0)
+
+ torch::with_no_grad({
+ for (i in seq_along(main_keys)) {
+ key <- main_keys[[i]]
+
+ if (flux_is_quant_key(key) &&
+ paste0(key, sib_suffix) %in% sib_keys) {
+ segments <- strsplit(key, ".", fixed = TRUE)[[1]]
+ parent <- .ltx23_walk_module(model, utils::head(segments, -2L))
+ leaf <- segments[length(segments) - 1L]
+ old <- .ltx23_walk_module(parent, leaf)
+ if (is.null(old)) {
+ unmapped <- c(unmapped, key)
+ next
+ }
+ w_shape <- old$weight$shape
+ has_bias <- !is.null(old$bias)
+ quant_mod <- if (format == "nf4") {
+ ltx23_nf4_linear(w_shape[1], w_shape[2], bias = has_bias)
+ } else {
+ ltx23_fp8_linear(w_shape[1], w_shape[2], bias = has_bias)
+ }
+ if (has_bias) {
+ # Adopt the original bias parameter; its checkpoint key
+ # loads through the pre-swap destination map
+ quant_mod$bias <- old$bias
+ }
+ if (format == "nf4") {
+ quant_mod$set_nf4_weight(
+ ckpt$handle$get_tensor(key),
+ ckpt$handle$get_tensor(paste0(key, sib_suffix))
+ )
+ } else {
+ quant_mod$set_fp8_weight(
+ ckpt$handle$get_tensor(key),
+ ckpt$handle$get_tensor(paste0(key, sib_suffix)),
+ pin = pin
+ )
+ }
+ do.call(`$<-`, list(parent, leaf, quant_mod))
+ filled <- c(filled, key)
+ } else {
+ dest <- dests[[key]]
+ if (is.null(dest)) {
+ unmapped <- c(unmapped, key)
+ next
+ }
+ dest$copy_(ckpt$handle$get_tensor(key))
+ filled <- c(filled, key)
+ }
+
+ if (i %% 100L == 0L) {
+ gc(verbose = FALSE)
+ if (verbose && i %% 500L == 0L) {
+ message(sprintf(" loaded %d/%d transformer tensors", i,
+ length(main_keys)))
+ }
+ }
+ }
+ })
+ gc(verbose = FALSE)
+
+ if (length(unmapped)) {
+ stop("FLUX ", format, " load: ", length(unmapped),
+ " unmapped keys, e.g. ",
+ paste(utils::head(unmapped, 3), collapse = ", "))
+ }
+ # Weight params replaced by quantized modules won't be "filled"
+ expected_missing <- flux_is_quant_key(names(dests))
+ unfilled <- setdiff(names(dests)[!expected_missing], filled)
+ if (length(unfilled)) {
+ stop("FLUX ", format, " load: ", length(unfilled),
+ " unfilled params, e.g. ",
+ paste(utils::head(unfilled, 3), collapse = ", "))
+ }
+
+ # NF4: everything (packed buffers included) onto the device.
+ # FP8: residents move; fp8 weights are plain fields and stay CPU-side.
+ model$to(device = device)
+ model$eval()
+ # Block intermediates are large at image resolutions; per-block gc
+ # keeps the quantized-linear temporaries bounded
+ options(diffuseR.block_gc = TRUE)
+ if (verbose) {
+ message("FLUX transformer ready (", format, ") on ", device)
+ }
+ model
+}
diff --git a/R/rope_flux.R b/R/rope_flux.R
new file mode 100644
index 0000000..d58813c
--- /dev/null
+++ b/R/rope_flux.R
@@ -0,0 +1,114 @@
+#' FLUX Rotary Positional Embeddings
+#'
+#' Fresh R port of the FLUX rotary positional embedding scheme from the
+#' diffusers reference implementation (Apache-2.0,
+#' src/diffusers/models/transformers/transformer_flux.py FluxPosEmbed and
+#' src/diffusers/models/embeddings.py get_1d_rotary_pos_embed /
+#' apply_rotary_emb). FLUX uses the interleaved adjacent-pair convention
+#' (use_real_unbind_dim = -1) with per-axis frequencies computed in
+#' float64 and applied in float32. Text tokens carry all-zero ids, so
+#' they receive the identity rotation.
+#'
+#' @name rope_flux
+NULL
+
+#' Build FLUX latent image position ids
+#'
+#' Position ids over the packed latent grid (latent height/2 x width/2).
+#' Channel 1 is always zero, channel 2 holds the row index, channel 3 the
+#' column index. Reference: FluxPipeline._prepare_latent_image_ids.
+#'
+#' @param height Integer. Packed grid height (latent height / 2).
+#' @param width Integer. Packed grid width (latent width / 2).
+#' @param device Device for the resulting tensor.
+#'
+#' @return Float tensor of shape [height * width, 3].
+#'
+#' @export
+flux_prepare_latent_image_ids <- function(height, width, device = "cpu") {
+ f32 <- torch::torch_float32()
+ ids <- torch::torch_zeros(height, width, 3L, dtype = f32, device = device)
+ # torch_arange has an inclusive end; end = n - 1 matches Python arange
+ rows <- torch::torch_arange(start = 0, end = height - 1, dtype = f32,
+ device = device)
+ cols <- torch::torch_arange(start = 0, end = width - 1, dtype = f32,
+ device = device)
+ ids[,, 2] <- ids[,, 2] + rows$unsqueeze(2L)
+ ids[,, 3] <- ids[,, 3] + cols$unsqueeze(1L)
+ ids$reshape(c(height * width, 3L))
+}
+
+#' Compute FLUX rotary frequencies from position ids
+#'
+#' Per-axis 1D rotary frequencies (interleaved-real convention), computed
+#' in float64 on CPU and concatenated over the axes. Reference:
+#' FluxPosEmbed with get_1d_rotary_pos_embed(repeat_interleave_real=TRUE,
+#' use_real=TRUE, freqs_dtype=float64).
+#'
+#' @param ids Tensor of shape [S, 3]: concatenated text ids (all zero)
+#' and image ids from \code{flux_prepare_latent_image_ids}.
+#' @param axes_dim Integer vector of per-axis rotary dims; must sum to
+#' the attention head dim. FLUX uses c(16, 56, 56).
+#' @param theta Numeric. RoPE base frequency.
+#'
+#' @return List of two tensors (cos, sin), each [S, sum(axes_dim)],
+#' float32, on the device of \code{ids}.
+#'
+#' @export
+flux_pos_embed <- function(ids, axes_dim = c(16L, 56L, 56L), theta = 10000) {
+ n_axes <- ids$shape[2]
+ device <- ids$device
+ f64 <- torch::torch_float64()
+ # Frequencies in float64 on CPU: Blackwell fp64 throughput is 1/64,
+ # and the tensors are tiny
+ pos <- ids$to(dtype = f64)$cpu()
+
+ cos_out <- vector("list", n_axes)
+ sin_out <- vector("list", n_axes)
+ for (i in seq_len(n_axes)) {
+ d <- axes_dim[i]
+ # freqs = 1 / theta^(seq(0, d - 2, by = 2) / d), length d / 2
+ exponents <- torch::torch_arange(start = 0, end = d - 2, step = 2,
+ dtype = f64)
+ freqs <- 1.0 / torch::torch_pow(theta, exponents / d)
+ # Outer product [S, d/2]
+ freqs <- pos[, i]$unsqueeze(2L) * freqs$unsqueeze(1L)
+ cos_out[[i]] <- freqs$cos()$repeat_interleave(2L, dim = 2L)
+ sin_out[[i]] <- freqs$sin()$repeat_interleave(2L, dim = 2L)
+ }
+
+ f32 <- torch::torch_float32()
+ list(
+ torch::torch_cat(cos_out, dim = -1L)$to(dtype = f32, device = device),
+ torch::torch_cat(sin_out, dim = -1L)$to(dtype = f32, device = device)
+ )
+}
+
+#' Apply FLUX rotary embeddings to a per-head tensor
+#'
+#' Rotates adjacent element pairs of the last dimension:
+#' \code{out = x * cos + rotate_half(x) * sin} with pairs interleaved
+#' (elements 1,2 form the first complex pair). Math in float32, result
+#' cast back to the input dtype. Reference: apply_rotary_emb with
+#' use_real_unbind_dim = -1.
+#'
+#' @param x Tensor of shape [B, H, S, D] (per-head layout).
+#' @param freqs List of two tensors (cos, sin), each [S, D], from
+#' \code{flux_pos_embed}.
+#'
+#' @return Tensor with the same shape and dtype as \code{x}.
+#'
+#' @export
+flux_apply_rotary_emb <- function(x, freqs) {
+ cos <- freqs[[1]]$unsqueeze(1L)$unsqueeze(1L) # [1, 1, S, D]
+ sin <- freqs[[2]]$unsqueeze(1L)$unsqueeze(1L)
+
+ pairs <- x$unflatten(4L, c(-1L, 2L)) # [B, H, S, D/2, 2]
+ x_real <- pairs[,,,, 1]
+ x_imag <- pairs[,,,, 2]
+ x_rotated <- torch::torch_stack(list(-x_imag, x_real), dim = -1L)$flatten(start_dim = 4L)
+
+ out <- x$to(dtype = torch::torch_float32()) * cos +
+ x_rotated$to(dtype = torch::torch_float32()) * sin
+ out$to(dtype = x$dtype)
+}
diff --git a/R/t5_text_encoder.R b/R/t5_text_encoder.R
new file mode 100644
index 0000000..5f3e5bb
--- /dev/null
+++ b/R/t5_text_encoder.R
@@ -0,0 +1,306 @@
+#' T5 Text Encoder (T5-v1.1)
+#'
+#' Fresh R port of the T5 encoder stack from HuggingFace transformers
+#' (Apache-2.0, src/transformers/models/t5/modeling_t5.py), as used by
+#' FLUX's second text encoder (T5-v1.1-XXL: 24 layers, d_model 4096,
+#' 64 heads x d_kv 64, gated-GELU FFN). Distinctives faithfully carried
+#' over: RMS layer norms (no mean subtraction), no biases anywhere, no
+#' 1/sqrt(d) attention scaling (folded into the weights), and a shared
+#' relative position bias computed once from block 1's embedding and
+#' added to every layer's attention logits. Module field names mirror
+#' the checkpoint keys (minus the \code{encoder.} prefix).
+#'
+#' FLUX passes no attention mask - padding tokens attend and are
+#' attended to - so none is implemented.
+#'
+#' @name t5_text_encoder
+NULL
+
+# Bucketed relative positions (bidirectional): half the buckets split by
+# sign, half of those exact small offsets, the rest log-spaced up to
+# max_distance. Reference: T5Attention._relative_position_bucket.
+.t5_relative_position_bucket <- function(relative_position,
+ num_buckets = 32L,
+ max_distance = 128L) {
+ num_buckets <- num_buckets %/% 2L
+ long <- torch::torch_long()
+ relative_buckets <- (relative_position > 0)$to(dtype = long)$mul(num_buckets)
+ relative_position <- torch::torch_abs(relative_position)
+
+ max_exact <- num_buckets %/% 2L
+ is_small <- relative_position < max_exact
+
+ rp_large <- relative_position$to(dtype = torch::torch_float32())$
+ div(max_exact)$log()$
+ div(log(max_distance / max_exact))$
+ mul(num_buckets - max_exact)$
+ to(dtype = long)$add(max_exact)
+ rp_large <- torch::torch_minimum(
+ rp_large,
+ torch::torch_full_like(rp_large, num_buckets - 1L)
+ )
+
+ relative_buckets + torch::torch_where(is_small, relative_position, rp_large)
+}
+
+# T5 self-attention: no scaling, no biases; the relative position bias
+# is added to the logits pre-softmax (softmax in float32)
+.t5_attention <- torch::nn_module(
+ "t5_attention",
+ initialize = function(d_model, d_kv, num_heads, has_relative_bias = FALSE,
+ num_buckets = 32L, max_distance = 128L) {
+ inner_dim <- num_heads * d_kv
+ self$num_heads <- num_heads
+ self$d_kv <- d_kv
+ self$num_buckets <- as.integer(num_buckets)
+ self$max_distance <- as.integer(max_distance)
+ self$q <- torch::nn_linear(d_model, inner_dim, bias = FALSE)
+ self$k <- torch::nn_linear(d_model, inner_dim, bias = FALSE)
+ self$v <- torch::nn_linear(d_model, inner_dim, bias = FALSE)
+ self$o <- torch::nn_linear(inner_dim, d_model, bias = FALSE)
+ if (has_relative_bias) {
+ self$relative_attention_bias <- torch::nn_embedding(num_buckets,
+ num_heads)
+ }
+},
+ compute_bias = function(seq_len, device) {
+ pos <- torch::torch_arange(start = 0, end = seq_len - 1,
+ dtype = torch::torch_long(), device = device)
+ # relative_position[i, j] = j - i
+ relative_position <- pos$unsqueeze(1L) - pos$unsqueeze(2L)
+ buckets <- .t5_relative_position_bucket(relative_position,
+ num_buckets = self$num_buckets,
+ max_distance = self$max_distance)
+ values <- self$relative_attention_bias(buckets + 1L) # [S, S, H]
+ values$permute(c(3L, 1L, 2L))$unsqueeze(1L) # [1, H, S, S]
+},
+ forward = function(x, position_bias) {
+ shape <- x$shape
+ b <- shape[1]
+ s <- shape[2]
+ per_head <- c(b, s, self$num_heads, self$d_kv)
+ q <- self$q(x)$view(per_head)$transpose(2L, 3L) # [B, H, S, dk]
+ k <- self$k(x)$view(per_head)$transpose(2L, 3L)
+ v <- self$v(x)$view(per_head)$transpose(2L, 3L)
+
+ scores <- torch::torch_matmul(q, k$transpose(-2L, -1L)) # no 1/sqrt(d)
+ scores <- scores + position_bias
+ attn <- torch::nnf_softmax(scores$to(dtype = torch::torch_float32()),
+ dim = -1L)$to(dtype = scores$dtype)
+ out <- torch::torch_matmul(attn, v)
+ out <- out$transpose(2L, 3L)$reshape(c(b, s, -1L))
+ self$o(out)
+}
+)
+
+# layer.0: pre-norm self-attention with residual
+.t5_self_attn_layer <- torch::nn_module(
+ "t5_self_attn_layer",
+ initialize = function(d_model, d_kv, num_heads, eps,
+ has_relative_bias = FALSE, num_buckets = 32L,
+ max_distance = 128L) {
+ self$SelfAttention <- .t5_attention(d_model, d_kv, num_heads,
+ has_relative_bias = has_relative_bias,
+ num_buckets = num_buckets,
+ max_distance = max_distance)
+ self$layer_norm <- ltx23_rms_norm(d_model, eps = eps)
+},
+ forward = function(x, position_bias) {
+ x + self$SelfAttention(self$layer_norm(x), position_bias)
+}
+)
+
+# layer.1: pre-norm gated-GELU feed-forward with residual
+.t5_ff_layer <- torch::nn_module(
+ "t5_ff_layer",
+ initialize = function(d_model, d_ff, eps) {
+ dense <- torch::nn_module(
+ "t5_dense_gated_act_dense",
+ initialize = function(d_model, d_ff) {
+ self$wi_0 <- torch::nn_linear(d_model, d_ff, bias = FALSE)
+ self$wi_1 <- torch::nn_linear(d_model, d_ff, bias = FALSE)
+ self$wo <- torch::nn_linear(d_ff, d_model, bias = FALSE)
+ },
+ forward = function(x) {
+ h <- torch::nnf_gelu(self$wi_0(x), approximate = "tanh") * self$wi_1(x)
+ self$wo(h)
+ }
+ )
+ self$DenseReluDense <- dense(d_model, d_ff)
+ self$layer_norm <- ltx23_rms_norm(d_model, eps = eps)
+},
+ forward = function(x) {
+ x + self$DenseReluDense(self$layer_norm(x))
+}
+)
+
+.t5_block <- torch::nn_module(
+ "t5_block",
+ initialize = function(d_model, d_kv, num_heads, d_ff, eps,
+ has_relative_bias = FALSE, num_buckets = 32L,
+ max_distance = 128L) {
+ self$layer <- torch::nn_module_list(list(
+ .t5_self_attn_layer(d_model, d_kv, num_heads, eps,
+ has_relative_bias = has_relative_bias,
+ num_buckets = num_buckets,
+ max_distance = max_distance),
+ .t5_ff_layer(d_model, d_ff, eps)
+ ))
+},
+ forward = function(x, position_bias) {
+ x <- self$layer[[1]](x, position_bias)
+ self$layer[[2]](x)
+}
+)
+
+#' T5 encoder stack
+#'
+#' Defaults are the T5-v1.1-XXL configuration used by FLUX.
+#'
+#' @param vocab_size,d_model,d_kv,num_heads,d_ff,num_layers Integers.
+#' @param relative_attention_num_buckets,relative_attention_max_distance
+#' Integers. Relative position bias shape.
+#' @param layer_norm_epsilon Numeric.
+#'
+#' @return Module whose forward(input_ids) (1-based ids [B, S]) returns
+#' the last hidden state [B, S, d_model].
+#'
+#' @export
+t5_encoder <- torch::nn_module(
+ "t5_encoder",
+ initialize = function(vocab_size = 32128L, d_model = 4096L, d_kv = 64L,
+ num_heads = 64L, d_ff = 10240L, num_layers = 24L,
+ relative_attention_num_buckets = 32L,
+ relative_attention_max_distance = 128L,
+ layer_norm_epsilon = 1e-6) {
+ self$shared <- torch::nn_embedding(vocab_size, d_model)
+ self$block <- torch::nn_module_list(
+ lapply(seq_len(num_layers), function(i) {
+ .t5_block(d_model, d_kv, num_heads, d_ff, layer_norm_epsilon,
+ has_relative_bias = (i == 1L),
+ num_buckets = relative_attention_num_buckets,
+ max_distance = relative_attention_max_distance)
+ })
+ )
+ self$final_layer_norm <- ltx23_rms_norm(d_model, eps = layer_norm_epsilon)
+},
+ forward = function(input_ids) {
+ x <- self$shared(input_ids)
+ # Bias comes from block 1 and is shared by every layer
+ position_bias <- self$block[[1]]$layer[[1]]$SelfAttention$compute_bias(
+ input_ids$shape[2], input_ids$device
+ )$to(dtype = x$dtype)
+ for (i in seq_along(self$block)) {
+ x <- self$block[[i]](x, position_bias)
+ }
+ self$final_layer_norm(x)
+}
+)
+
+# Encoder constructor arguments from a transformers config.json
+.t5_encoder_args <- function(config) {
+ if (is.null(config)) {
+ return(list())
+ }
+ ffp <- config$feed_forward_proj %||% "gated-gelu"
+ if (!startsWith(ffp, "gated")) {
+ stop("Only gated feed-forward T5 (v1.1) is supported, got: ", ffp)
+ }
+ args <- list(
+ vocab_size = config$vocab_size,
+ d_model = config$d_model,
+ d_kv = config$d_kv,
+ num_heads = config$num_heads,
+ d_ff = config$d_ff,
+ num_layers = config$num_layers,
+ relative_attention_num_buckets = config$relative_attention_num_buckets,
+ relative_attention_max_distance = config$relative_attention_max_distance
+ )
+ args <- Filter(function(x) !is.null(x) && length(x) > 0L, args)
+ args <- lapply(args, as.integer)
+ eps <- config$layer_norm_epsilon
+ if (!is.null(eps) && length(eps) == 1L) {
+ args$layer_norm_epsilon <- as.numeric(eps)
+ }
+ args
+}
+
+#' Load a T5 encoder from a transformers directory
+#'
+#' Streams the (possibly sharded) safetensors weights into
+#' \code{\link{t5_encoder}}, stripping the \code{encoder.} key prefix
+#' and aliasing \code{embed_tokens} to the shared embedding.
+#'
+#' @param model_path Directory with \code{config.json} and
+#' \code{model*.safetensors} (FLUX.1-schnell's \code{text_encoder_2}).
+#' @param device Character. Target device.
+#' @param dtype Character. "float32" (CPU default; T5 overflows in
+#' float16) or "bfloat16".
+#' @param verbose Logical.
+#' @param ... Overrides for \code{\link{t5_encoder}} arguments.
+#'
+#' @return The loaded \code{t5_encoder} in eval mode.
+#'
+#' @export
+load_t5_text_encoder <- function(model_path, device = "cpu",
+ dtype = "float32", verbose = TRUE, ...) {
+ model_path <- path.expand(model_path)
+ config <- NULL
+ config_path <- file.path(model_path, "config.json")
+ if (file.exists(config_path)) {
+ config <- jsonlite::fromJSON(config_path, simplifyVector = TRUE)
+ }
+
+ args <- utils::modifyList(.t5_encoder_args(config), list(...))
+ model <- do.call(t5_encoder, args)
+ model$to(dtype = .flux_dtype(dtype))
+
+ opened <- .flux_open_sharded_dir(model_path, "model")
+ ckpt <- structure(
+ list(handle = opened$handle, keys = opened$keys, version = NULL,
+ config = config, path = model_path),
+ class = "ltx23_checkpoint"
+ )
+
+ map_key <- function(key) {
+ key <- sub("^encoder\\.", "", key)
+ if (key == "embed_tokens.weight") {
+ key <- "shared.weight"
+ }
+ key
+ }
+ res <- ltx23_load_group(ckpt, ckpt$keys, model, map_key = map_key,
+ verbose = verbose)
+ if (length(res$unmapped) || length(res$unfilled)) {
+ stop("T5 encoder load: ", length(res$unmapped), " unmapped keys, ",
+ length(res$unfilled), " unfilled params")
+ }
+
+ model$to(device = device)
+ model$eval()
+ model
+}
+
+#' Encode prompts with the T5 encoder
+#'
+#' Tokenizes with \code{\link{encode_unigram}} (right padding to
+#' \code{max_sequence_length}) and runs the encoder. Matching the FLUX
+#' reference pipeline, no attention mask is used.
+#'
+#' @param prompts Character vector.
+#' @param model A \code{\link{t5_encoder}}.
+#' @param tokenizer A \code{\link{unigram_tokenizer}}.
+#' @param max_sequence_length Integer. Fixed token length (schnell: 256).
+#' @param device Device for the input ids (defaults to the model's).
+#'
+#' @return Tensor [length(prompts), max_sequence_length, d_model].
+#'
+#' @export
+encode_with_t5 <- function(prompts, model, tokenizer,
+ max_sequence_length = 256L, device = NULL) {
+ enc <- encode_unigram(tokenizer, prompts, max_length = max_sequence_length)
+ device <- device %||% model$shared$weight$device
+ ids <- torch::torch_tensor(enc$input_ids + 1L,
+ dtype = torch::torch_long(), device = device)
+ torch::with_no_grad(model(ids))
+}
diff --git a/R/text_encoder.R b/R/text_encoder.R
index 6cd04e3..60687fc 100644
--- a/R/text_encoder.R
+++ b/R/text_encoder.R
@@ -157,6 +157,8 @@ CLIPTransformerBlock <- torch::nn_module(
#' @param mlp_dim MLP hidden dimension
#' @param apply_final_ln Whether to apply final layer norm (default TRUE).
#' Set to FALSE to match TorchScript exports that don't include final LN.
+#' @param gelu_type GELU variant: "tanh" (matches the TorchScript exports),
+#' "quick" (HF CLIP ViT-L, used by FLUX), or "exact"
#'
#' @return An nn_module representing the text encoder
#' @export
@@ -170,7 +172,8 @@ text_encoder_native <- torch::nn_module(
num_layers = 12,
num_heads = 12,
mlp_dim = 3072,
- apply_final_ln = TRUE
+ apply_final_ln = TRUE,
+ gelu_type = "tanh"
) {
self$context_length <- context_length
self$embed_dim <- embed_dim
@@ -183,12 +186,11 @@ text_encoder_native <- torch::nn_module(
torch::torch_zeros(context_length, embed_dim)
)
- # Transformer blocks - tanh GELU approximation matches TorchScript export best
self$transformer_blocks <- torch::nn_module_list()
for (i in seq_len(num_layers)) {
self$transformer_blocks$append(
CLIPTransformerBlock(embed_dim, num_heads, mlp_dim,
- gelu_type = "tanh")
+ gelu_type = gelu_type)
)
}
@@ -290,6 +292,100 @@ detect_text_encoder_architecture <- function(torchscript_path) {
)
}
+#' Pooled CLIP output at the EOS position
+#'
+#' The HF CLIPTextModel pooler_output: the final-layer-norm hidden state
+#' at the EOS token position, located by argmax over the token ids (EOS
+#' is the highest id in the CLIP vocab, and causal attention makes any
+#' padding after it irrelevant). No text projection is applied - this is
+#' what FLUX uses as pooled_projections.
+#'
+#' @param hidden_states Final-LN hidden states [B, S, D] from
+#' \code{\link{text_encoder_native}} (with \code{apply_final_ln = TRUE}).
+#' @param input_ids Token ids [B, S] (0-based, as fed to the encoder).
+#'
+#' @return Tensor [B, D].
+#'
+#' @export
+clip_pooled_output <- function(hidden_states, input_ids) {
+ input_ids <- input_ids$to(device = hidden_states$device)
+ eos_indices <- torch::torch_argmax(input_ids, dim = 2L, keepdim = TRUE)
+ hidden_states$gather(
+ dim = 2L,
+ index = eos_indices$unsqueeze(-1L)$expand(c(-1L, -1L,
+ hidden_states$shape[3]))
+ )$squeeze(2L)
+}
+
+#' Load HF safetensors weights into the native CLIP text encoder
+#'
+#' Loads a HuggingFace CLIPTextModel \code{model.safetensors} (e.g.
+#' FLUX.1-schnell's \code{text_encoder}) into
+#' \code{\link{text_encoder_native}}, reusing the TorchScript key remaps
+#' minus the export prefixes.
+#'
+#' @param native_encoder Native text encoder module
+#' @param path Path to model.safetensors (or a directory containing it)
+#' @param verbose Print loading progress
+#'
+#' @return The native encoder with loaded weights (invisibly)
+#' @export
+load_text_encoder_safetensors <- function(native_encoder, path,
+ verbose = TRUE) {
+ path <- path.expand(path)
+ if (dir.exists(path)) {
+ path <- file.path(path, "model.safetensors")
+ }
+ handle <- safetensors::safetensors$new(path, framework = "torch")
+ keys <- setdiff(handle$keys(), "__metadata__")
+ # Non-parameter buffers in some exports
+ keys <- keys[!endsWith(keys, "position_ids")]
+
+ remap_key <- function(key) {
+ key <- sub("^text_model\\.", "", key)
+ key <- sub("^embeddings\\.token_embedding\\.", "token_embedding.", key)
+ key <- sub("^embeddings\\.position_embedding\\.weight$",
+ "position_embedding", key)
+ key <- gsub("^encoder\\.layers\\.", "transformer_blocks.", key)
+ key <- gsub("\\.self_attn\\.", ".attention.", key)
+ key <- gsub("\\.layer_norm1\\.", ".layernorm_1.", key)
+ key <- gsub("\\.layer_norm2\\.", ".layernorm_2.", key)
+ key
+ }
+
+ dests <- native_encoder$named_parameters()
+ filled <- character(0)
+ unmapped <- character(0)
+ torch::with_no_grad({
+ for (key in keys) {
+ native_name <- remap_key(key)
+ dest <- dests[[native_name]]
+ if (is.null(dest)) {
+ unmapped <- c(unmapped, key)
+ next
+ }
+ dest$copy_(handle$get_tensor(key))
+ filled <- c(filled, native_name)
+ }
+ })
+
+ unfilled <- setdiff(names(dests), filled)
+ if (length(unmapped)) {
+ stop("CLIP safetensors load: ", length(unmapped),
+ " unmapped keys, e.g. ",
+ paste(utils::head(unmapped, 3), collapse = ", "))
+ }
+ if (length(unfilled)) {
+ stop("CLIP safetensors load: ", length(unfilled),
+ " unfilled params, e.g. ",
+ paste(utils::head(unfilled, 3), collapse = ", "))
+ }
+ if (verbose) {
+ message("Loaded ", length(filled), " CLIP parameters from ", path)
+ }
+ invisible(native_encoder)
+}
+
#' Load weights from TorchScript text encoder into native encoder
#'
#' @param native_encoder Native text encoder module
diff --git a/R/tokenizer_unigram.R b/R/tokenizer_unigram.R
new file mode 100644
index 0000000..677afff
--- /dev/null
+++ b/R/tokenizer_unigram.R
@@ -0,0 +1,270 @@
+#' SentencePiece Unigram Tokenizer
+#'
+#' Pure R implementation of HuggingFace tokenizer.json files with a
+#' Unigram model (SentencePiece), as used by T5 - FLUX's second text
+#' encoder. Segmentation is Viterbi best-path over the vocab log
+#' probabilities (Kudo 2018, arXiv:1804.10959). The normalizer and
+#' Metaspace pre-tokenizer settings are read from the file.
+#'
+#' Limitation: the Precompiled charsmap normalizer (NFKC-style unicode
+#' mapping) is approximated by control-whitespace substitution only;
+#' ASCII and common latin text tokenizes identically to the reference,
+#' exotic unicode may differ.
+#'
+#' @name tokenizer_unigram
+NULL
+
+#' Load a Unigram tokenizer from tokenizer.json
+#'
+#' @param tokenizer_path Path to a HuggingFace tokenizer.json with a
+#' Unigram model, or a directory containing one.
+#'
+#' @return A \code{unigram_tokenizer} object.
+#'
+#' @export
+unigram_tokenizer <- function(tokenizer_path) {
+ path <- path.expand(tokenizer_path)
+ if (dir.exists(path)) {
+ path <- file.path(path, "tokenizer.json")
+ }
+ if (!file.exists(path)) {
+ stop("tokenizer.json not found: ", path)
+ }
+
+ tj <- jsonlite::fromJSON(path, simplifyVector = FALSE)
+ model <- tj$model
+ if (is.null(model) || !identical(model$type, "Unigram")) {
+ stop("Only Unigram tokenizers are supported (got ",
+ model$type %||% "none", "); use bpe_tokenizer() for BPE.")
+ }
+
+ pieces <- vapply(model$vocab, function(p) p[[1]], character(1))
+ scores <- vapply(model$vocab, function(p) as.numeric(p[[2]]), numeric(1))
+ ids0 <- seq_along(pieces) - 1L
+
+ vocab_env <- new.env(parent = emptyenv(), size = length(pieces))
+ for (i in seq_along(pieces)) {
+ assign(pieces[i], c(ids0[i], scores[i]), envir = vocab_env)
+ }
+
+ # Normalizer settings (Sequence or single normalizer)
+ norms <- tj$normalizer
+ if (!is.null(norms) && identical(norms$type, "Sequence")) {
+ norms <- norms$normalizers
+ } else if (!is.null(norms)) {
+ norms <- list(norms)
+ } else {
+ norms <- list()
+ }
+ strip_right <- FALSE
+ space_collapse <- NULL
+ for (n in norms) {
+ if (identical(n$type, "Strip") && isTRUE(n$strip_right)) {
+ strip_right <- TRUE
+ }
+ if (identical(n$type, "Replace") &&
+ identical(n$pattern$Regex, " {2,}")) {
+ space_collapse <- n$content
+ }
+ }
+
+ # Metaspace pre-tokenizer (possibly inside a Sequence)
+ pre <- tj$pre_tokenizer
+ pres <- if (!is.null(pre) && identical(pre$type, "Sequence")) {
+ pre$pretokenizers
+ } else if (!is.null(pre)) {
+ list(pre)
+ } else {
+ list()
+ }
+ prepend_scheme <- "never"
+ replacement <- "\u2581"
+ for (p in pres) {
+ if (identical(p$type, "Metaspace")) {
+ replacement <- p$replacement %||% "\u2581"
+ prepend_scheme <- p$prepend_scheme %||%
+ (if (isTRUE(p$add_prefix_space)) "always" else "never")
+ }
+ }
+
+ unk_id <- as.integer(model$unk_id %||% 2L)
+ eos <- get0("", envir = vocab_env)
+ pad <- get0("", envir = vocab_env)
+
+ structure(
+ list(
+ vocab = vocab_env,
+ n_pieces = length(pieces),
+ max_piece_chars = max(nchar(pieces)),
+ unk_id = unk_id,
+ unk_score = min(scores) - 10.0,
+ eos_id = if (!is.null(eos)) as.integer(eos[1]) else 1L,
+ pad_id = if (!is.null(pad)) as.integer(pad[1]) else 0L,
+ strip_right = strip_right,
+ space_collapse = space_collapse,
+ prepend_scheme = prepend_scheme,
+ replacement = replacement,
+ path = path
+ ),
+ class = "unigram_tokenizer"
+ )
+}
+
+#' @export
+print.unigram_tokenizer <- function(x, ...) {
+ cat("\n")
+ cat(" pieces: ", x$n_pieces, "\n")
+ cat(" unk/eos/pad:", x$unk_id, x$eos_id, x$pad_id, "\n")
+ cat(" path: ", x$path, "\n")
+ invisible(x)
+}
+
+# Viterbi best-path segmentation of one pre-token (0-based ids)
+.unigram_viterbi <- function(word, tokenizer) {
+ n <- nchar(word)
+ if (n == 0L) {
+ return(integer(0))
+ }
+ vocab <- tokenizer$vocab
+ max_len <- tokenizer$max_piece_chars
+
+ best <- c(0, rep(-Inf, n))
+ back_len <- integer(n)
+ back_id <- integer(n)
+
+ for (end in seq_len(n)) {
+ for (len in seq_len(min(max_len, end))) {
+ start <- end - len + 1L
+ piece <- substr(word, start, end)
+ entry <- get0(piece, envir = vocab)
+ if (is.null(entry)) {
+ if (len > 1L) {
+ next
+ }
+ id <- tokenizer$unk_id
+ score <- tokenizer$unk_score
+ } else {
+ id <- entry[1]
+ score <- entry[2]
+ }
+ cand <- best[start] + score
+ if (cand > best[end + 1L]) {
+ best[end + 1L] <- cand
+ back_len[end] <- len
+ back_id[end] <- id
+ }
+ }
+ }
+
+ ids <- integer(0)
+ pos <- n
+ while (pos > 0L) {
+ ids <- c(back_id[pos], ids)
+ pos <- pos - back_len[pos]
+ }
+ ids
+}
+
+#' Encode text with a Unigram tokenizer
+#'
+#' Normalizes (strip-right, multi-space collapse, control whitespace to
+#' space), applies the Metaspace pre-tokenizer, segments each pre-token
+#' by Viterbi over the Unigram scores, fuses consecutive unknowns, and
+#' appends EOS. T5 semantics: right padding with \code{} (id 0),
+#' truncation to \code{max_length - 1} before the EOS.
+#'
+#' @param tokenizer A \code{\link{unigram_tokenizer}}.
+#' @param texts Character vector of prompts.
+#' @param max_length Integer. Fixed sequence length (NULL for no
+#' truncation/padding).
+#' @param add_eos Logical. Append the EOS token.
+#' @param pad Logical. Right-pad to \code{max_length}.
+#'
+#' @return List with \code{input_ids} and \code{attention_mask}, each an
+#' integer matrix [length(texts), max_length] (or ragged lists when
+#' \code{max_length} is NULL). Ids are 0-based (HuggingFace
+#' convention); add 1 for R torch embedding lookups.
+#'
+#' @export
+encode_unigram <- function(tokenizer, texts, max_length = 256L,
+ add_eos = TRUE, pad = TRUE) {
+ stopifnot(inherits(tokenizer, "unigram_tokenizer"))
+ rep_char <- tokenizer$replacement
+
+ encode_one <- function(text) {
+ # Control whitespace to space (charsmap approximation), then the
+ # file's normalizer chain
+ text <- enc2utf8(text)
+ text <- gsub("[\t\n\r\f\v]", " ", text)
+ if (tokenizer$strip_right) {
+ text <- sub("[ ]+$", "", text)
+ }
+ if (!is.null(tokenizer$space_collapse)) {
+ text <- gsub(" {2,}", tokenizer$space_collapse, text)
+ }
+ # Metaspace: spaces to the replacement, optional prefix.
+ # Empty input yields no pre-tokens (before any prepend).
+ text <- gsub(" ", rep_char, text, fixed = TRUE)
+ if (nchar(text) == 0L) {
+ return(integer(0))
+ }
+ if (tokenizer$prepend_scheme == "always" &&
+ !startsWith(text, rep_char)) {
+ text <- paste0(rep_char, text)
+ }
+ # Split before each replacement char, keeping it attached.
+ # (strsplit with a zero-width lookahead detaches the char, so
+ # mark boundaries with a control byte instead.)
+ marked <- gsub(rep_char, paste0("\u0001", rep_char), text, fixed = TRUE)
+ words <- strsplit(marked, "\u0001", fixed = TRUE)[[1]]
+ words <- words[nzchar(words)]
+
+ ids <- unlist(lapply(words, .unigram_viterbi, tokenizer = tokenizer))
+ ids <- as.integer(ids %||% integer(0))
+
+ # Fuse consecutive unknowns (HF fuse_unk)
+ if (length(ids) > 1L) {
+ is_unk <- ids == tokenizer$unk_id
+ drop <- is_unk & c(FALSE, is_unk[-length(is_unk)])
+ ids <- ids[!drop]
+ }
+ ids
+ }
+
+ all_ids <- lapply(as.character(texts), encode_one)
+
+ if (add_eos) {
+ if (is.null(max_length)) {
+ keep <- Inf
+ } else {
+ keep <- max_length - 1L
+ }
+ all_ids <- lapply(all_ids, function(ids) {
+ if (length(ids) > keep) {
+ ids <- ids[seq_len(keep)]
+ }
+ c(ids, tokenizer$eos_id)
+ })
+ } else if (!is.null(max_length)) {
+ all_ids <- lapply(all_ids, function(ids) {
+ if (length(ids) > max_length) ids[seq_len(max_length)] else ids
+ })
+ }
+
+ if (is.null(max_length) || !pad) {
+ masks <- lapply(all_ids, function(ids) rep(1L, length(ids)))
+ return(list(input_ids = all_ids, attention_mask = masks))
+ }
+
+ n <- length(all_ids)
+ input_ids <- matrix(tokenizer$pad_id, nrow = n, ncol = max_length)
+ attention_mask <- matrix(0L, nrow = n, ncol = max_length)
+ for (i in seq_len(n)) {
+ len <- length(all_ids[[i]])
+ if (len > 0L) {
+ input_ids[i, seq_len(len)] <- all_ids[[i]]
+ attention_mask[i, seq_len(len)] <- 1L
+ }
+ }
+ list(input_ids = input_ids, attention_mask = attention_mask)
+}
diff --git a/R/txt2img.R b/R/txt2img.R
index 445f3dc..476e282 100644
--- a/R/txt2img.R
+++ b/R/txt2img.R
@@ -11,11 +11,12 @@
#' \dontrun{
#' img <- txt2img("a cat wearing sunglasses in space", device = "cuda")
#' }
-txt2img <- function(prompt, model_name = c("sd21", "sdxl"), ...) {
+txt2img <- function(prompt, model_name = c("sd21", "sdxl", "flux1"), ...) {
switch(model_name,
# "sd15" = txt2img_sd15(prompt, ...),
"sd21" = txt2img_sd21(prompt, ...),
"sdxl" = txt2img_sdxl(prompt, ...),
+ "flux1" = txt2img_flux(prompt, ...),
# "sd3" = txt2img_sd3(prompt, ...),
stop("Unsupported model: ", model_name)
)
diff --git a/R/txt2img_flux.R b/R/txt2img_flux.R
new file mode 100644
index 0000000..1bf9282
--- /dev/null
+++ b/R/txt2img_flux.R
@@ -0,0 +1,437 @@
+#' FLUX.1 Text-to-Image Pipeline
+#'
+#' FLUX latent packing helpers and (in later phases) the schnell
+#' text-to-image pipeline. Ported from the diffusers reference
+#' implementation (Apache-2.0, src/diffusers/pipelines/flux/
+#' pipeline_flux.py).
+#'
+#' @name txt2img_flux
+NULL
+
+#' Pack FLUX latents into a patch sequence
+#'
+#' Packs a [B, C, H, W] latent into 2x2 patches, giving a sequence
+#' [B, (H/2) * (W/2), C * 4]. Reference: FluxPipeline._pack_latents.
+#'
+#' @param latents Tensor of shape [B, C, H, W]; H and W must be even.
+#'
+#' @return Tensor of shape [B, (H/2) * (W/2), C * 4].
+#'
+#' @export
+flux_pack_latents <- function(latents) {
+ shape <- latents$shape
+ b <- shape[1]
+ ch <- shape[2]
+ h <- shape[3]
+ w <- shape[4]
+ latents <- latents$view(c(b, ch, h %/% 2L, 2L, w %/% 2L, 2L))
+ # Python permute (0, 2, 4, 1, 3, 5), 1-indexed here
+ latents <- latents$permute(c(1L, 3L, 5L, 2L, 4L, 6L))
+ latents$reshape(c(b, (h %/% 2L) * (w %/% 2L), ch * 4L))
+}
+
+#' Unpack a FLUX patch sequence back into latents
+#'
+#' Inverse of \code{flux_pack_latents}. Height and width are the target
+#' image dimensions in pixels; the latent grid is derived via the VAE
+#' scale factor and the 2x2 patch size. Reference:
+#' FluxPipeline._unpack_latents.
+#'
+#' @param latents Tensor of shape [B, S, C_packed].
+#' @param height,width Integers. Image height/width in pixels.
+#' @param vae_scale_factor Integer. Spatial downsampling of the VAE (8).
+#'
+#' @return Tensor of shape [B, C_packed / 4, height / 8, width / 8].
+#'
+#' @export
+flux_unpack_latents <- function(latents, height, width, vae_scale_factor = 8L) {
+ shape <- latents$shape
+ b <- shape[1]
+ ch <- shape[3]
+ h <- 2L * (as.integer(height) %/% (vae_scale_factor * 2L))
+ w <- 2L * (as.integer(width) %/% (vae_scale_factor * 2L))
+ latents <- latents$view(c(b, h %/% 2L, w %/% 2L, ch %/% 4L, 2L, 2L))
+ # Python permute (0, 3, 1, 4, 2, 5), 1-indexed here
+ latents <- latents$permute(c(1L, 4L, 2L, 5L, 3L, 6L))
+ latents$reshape(c(b, ch %/% 4L, h, w))
+}
+
+# Resolve a FLUX.1-schnell support file from the HuggingFace cache
+.flux1_cached <- function(file) {
+ if (!requireNamespace("hfhub", quietly = TRUE)) {
+ stop("The hfhub package is required to locate model files.")
+ }
+ tryCatch(
+ hfhub::hub_download(.flux1_repo, file, local_files_only = TRUE),
+ error = function(e) {
+ stop("Missing ", file, " in the HuggingFace cache; ",
+ "run download_flux1() first.", call. = FALSE)
+ }
+ )
+}
+
+#' Load the FLUX.1-schnell pipeline
+#'
+#' Loads the quantized transformer artifact plus the VAE decoder, CLIP
+#' and T5 text encoders, tokenizer, and scheduler config (from the
+#' HuggingFace cache populated by \code{\link{download_flux1}}).
+#' Components load to the CPU when \code{phase_offload} is on and move
+#' to the GPU only for their phase of the generation.
+#'
+#' @param model_dir Quantized artifact directory (default: the
+#' \code{download_flux1} location for \code{precision}), or a raw
+#' diffusers transformer directory for full-precision loading.
+#' @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 attn_chunk Integer or NULL. Attention query-chunk override.
+#' @param phase_offload Logical. One GPU tenant per phase.
+#' @param verbose Logical.
+#'
+#' @return A \code{flux_pipeline} list.
+#'
+#' @export
+flux_load_pipeline <- function(model_dir = NULL, device = "cuda",
+ precision = NULL, text_device = "cpu",
+ attn_chunk = NULL, phase_offload = TRUE,
+ verbose = TRUE) {
+ profile <- flux_memory_profile()
+ if (is.null(precision)) {
+ precision <- profile$precision
+ }
+ if (is.null(attn_chunk)) {
+ attn_chunk <- profile$attn_chunk
+ }
+ if (is.null(model_dir)) {
+ model_dir <- file.path(tools::R_user_dir("diffuseR", "data"),
+ paste0("flux1-schnell-", precision))
+ }
+
+ ckpt <- if (file.exists(file.path(model_dir, "manifest.json"))) {
+ flux_open_quantized(model_dir)
+ } else {
+ flux_open_checkpoint(model_dir)
+ }
+
+ if (!nzchar(Sys.getenv("PYTORCH_CUDA_ALLOC_CONF"))) {
+ # Must be set before the first CUDA allocation (see
+ # ltx23_load_pipeline for the per-format rationale)
+ conf <- if (identical(ckpt$format, "nf4")) {
+ "backend:native"
+ } else {
+ "expandable_segments:True"
+ }
+ Sys.setenv(PYTORCH_CUDA_ALLOC_CONF = conf)
+ }
+ if (device == "cuda") {
+ if (identical(ckpt$format, "nf4")) {
+ footprint <- 8
+ } else {
+ footprint <- 4
+ }
+ ltx23_tune_gc(footprint_gb = footprint)
+ }
+
+ if (phase_offload) {
+ component_device <- "cpu"
+ } else {
+ component_device <- device
+ }
+
+ pipe <- list(
+ format = ckpt$format %||% "full",
+ device = device,
+ text_device = text_device,
+ phase_offload = phase_offload,
+ attn_chunk = if (is.null(attn_chunk)) NULL else as.integer(attn_chunk),
+ config = ckpt$config
+ )
+
+ if (verbose) {
+ message("Loading transformer (", pipe$format, ")...")
+ }
+ pipe$transformer <- flux_load_transformer(
+ ckpt, device = component_device,
+ dtype = if (device == "cpu") "float32" else "bfloat16",
+ pin = device == "cuda",
+ verbose = verbose
+ )
+
+ vae_config <- jsonlite::fromJSON(.flux1_cached("vae/config.json"))
+ pipe$vae_scaling_factor <- vae_config$scaling_factor %||% 0.3611
+ pipe$vae_shift_factor <- vae_config$shift_factor %||% 0.1159
+ if (verbose) {
+ message("Loading VAE decoder...")
+ }
+ dec <- vae_decoder_native(
+ latent_channels = as.integer(vae_config$latent_channels %||% 16L)
+ )
+ load_decoder_safetensors(
+ dec, .flux1_cached("vae/diffusion_pytorch_model.safetensors"),
+ verbose = verbose
+ )
+ dec$to(device = component_device)
+ dec$eval()
+ pipe$decoder <- dec
+
+ if (verbose) {
+ message("Loading CLIP text encoder...")
+ }
+ clip <- text_encoder_native(gelu_type = "quick")
+ load_text_encoder_safetensors(
+ clip, .flux1_cached("text_encoder/model.safetensors"),
+ verbose = verbose
+ )
+ clip$to(device = text_device)
+ clip$eval()
+ pipe$text_encoder <- clip
+
+ if (verbose) {
+ message("Loading T5 text encoder...")
+ }
+ t5_dir <- dirname(.flux1_cached("text_encoder_2/config.json"))
+ pipe$text_encoder2 <- load_t5_text_encoder(
+ t5_dir, device = text_device,
+ dtype = if (text_device == "cpu") "float32" else "bfloat16",
+ verbose = verbose
+ )
+ pipe$tokenizer2 <- unigram_tokenizer(.flux1_cached("tokenizer_2/tokenizer.json"))
+
+ sched_cfg <- tryCatch(
+ jsonlite::fromJSON(.flux1_cached("scheduler/scheduler_config.json")),
+ error = function(e) NULL
+ )
+ pipe$scheduler_shift <- sched_cfg$shift %||% 1.0
+
+ structure(pipe, class = "flux_pipeline")
+}
+
+# Flow-matching Euler loop over the packed latent sequence. Latents stay
+# float32; the transformer runs in compute_dtype. schnell is CFG-free:
+# one forward per step.
+.flux_denoise <- function(transformer, latents, schedule, prompt_embeds,
+ pooled_prompt_embeds, image_rotary_emb,
+ compute_dtype, chunk_size = NULL, verbose = TRUE) {
+ timesteps <- as.numeric(schedule$timesteps$cpu())
+ n <- length(timesteps)
+ pb <- if (verbose) {
+ utils::txtProgressBar(min = 0, max = n, style = 3)
+ } else {
+ NULL
+ }
+ f32 <- torch::torch_float32()
+
+ torch::with_no_grad({
+ for (i in seq_len(n)) {
+ t <- timesteps[[i]]
+ # The transformer takes sigma-space time (it rescales by 1000
+ # internally, matching the reference pipeline's t/1000)
+ t_model <- torch::torch_tensor(t / 1000, dtype = compute_dtype,
+ device = latents$device)$reshape(1L)
+
+ noise_pred <- transformer(
+ hidden_states = latents$to(dtype = compute_dtype),
+ encoder_hidden_states = prompt_embeds,
+ pooled_projections = pooled_prompt_embeds,
+ timestep = t_model,
+ image_rotary_emb = image_rotary_emb,
+ chunk_size = chunk_size
+ )
+
+ step <- flowmatch_scheduler_step(
+ noise_pred$to(dtype = f32), t, latents, schedule
+ )
+ latents <- step$prev_sample
+ schedule <- step$schedule
+ rm(noise_pred, step)
+ if (!is.null(pb)) {
+ utils::setTxtProgressBar(pb, i)
+ }
+ }
+ })
+ if (!is.null(pb)) {
+ close(pb)
+ }
+ latents
+}
+
+#' Generate an image with FLUX.1-schnell
+#'
+#' 4-step distilled text-to-image generation (no classifier-free
+#' guidance): T5 + CLIP prompt encoding, flow-matching Euler denoising
+#' over the packed latent sequence, and 16-channel VAE decode. With
+#' phase offloading each component is the sole GPU tenant for its phase.
+#'
+#' @param prompt Character. The prompt.
+#' @param pipeline A \code{flux_pipeline} from
+#' \code{\link{flux_load_pipeline}}; NULL loads one (passing
+#' \code{...} through).
+#' @param width,height Integers, divisible by 16.
+#' @param num_inference_steps Integer. Denoising steps (schnell: 4).
+#' @param max_sequence_length Integer. T5 token length (schnell: 256).
+#' @param seed Integer or NULL. Initial latents are drawn on the CPU, so
+#' a seed matches a Python diffusers run with a CPU generator.
+#' @param prompt_embeds,pooled_prompt_embeds Optional precomputed text
+#' embeddings (skip the text encoders).
+#' @param save_file Logical. Write a PNG.
+#' @param filename Output path (default derived from the prompt).
+#' @param verbose Logical.
+#' @param ... Passed to \code{\link{flux_load_pipeline}} when
+#' \code{pipeline} is NULL.
+#'
+#' @return Invisibly, \code{list(image, metadata)} where \code{image} is
+#' an [H, W, 3] array in [0, 1].
+#'
+#' @export
+txt2img_flux <- function(prompt, pipeline = NULL, width = 1024L,
+ height = 1024L, num_inference_steps = 4L,
+ max_sequence_length = 256L, seed = NULL,
+ prompt_embeds = NULL, pooled_prompt_embeds = NULL,
+ save_file = TRUE, filename = NULL, verbose = TRUE,
+ ...) {
+ if (is.null(pipeline)) {
+ pipeline <- flux_load_pipeline(..., verbose = verbose)
+ }
+ device <- pipeline$device
+ width <- as.integer(width)
+ height <- as.integer(height)
+ if (width %% 16L != 0L || height %% 16L != 0L) {
+ stop("width and height must be divisible by 16")
+ }
+
+ f32 <- torch::torch_float32()
+ compute_dtype <- if (device == "cpu") {
+ f32
+ } else {
+ torch::torch_bfloat16()
+ }
+
+ phase_offload <- isTRUE(pipeline$phase_offload) && device != "cpu"
+ onload <- function(module) {
+ if (phase_offload) {
+ module$to(device = device)
+ }
+ module
+ }
+ offload <- function(module) {
+ if (phase_offload) {
+ module$to(device = "cpu")
+ clear_vram()
+ }
+ invisible(module)
+ }
+
+ t0 <- Sys.time()
+
+ # --- Phase 1: text encoding ------------------------------------------------
+ torch::with_no_grad({
+ if (is.null(prompt_embeds)) {
+ if (verbose) {
+ message("Encoding prompt (T5)...")
+ }
+ prompt_embeds <- encode_with_t5(prompt, pipeline$text_encoder2,
+ pipeline$tokenizer2, max_sequence_length = max_sequence_length)
+ }
+ if (is.null(pooled_prompt_embeds)) {
+ 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)
+ }
+ })
+ prompt_embeds <- prompt_embeds$to(device = device, dtype = compute_dtype)
+ pooled_prompt_embeds <- pooled_prompt_embeds$to(device = device,
+ dtype = compute_dtype)
+ txt_len <- prompt_embeds$shape[2]
+
+ # --- Phase 2: latents, rotary embeddings, schedule --------------------------
+ lat_ch <- as.integer((pipeline$config$in_channels %||% 64L) %/% 4L)
+ h_lat <- 2L * (height %/% 16L)
+ w_lat <- 2L * (width %/% 16L)
+ if (!is.null(seed)) {
+ torch::torch_manual_seed(seed)
+ }
+ # Drawn on the CPU in the unpacked diffusers shape for seed parity
+ latents <- torch::torch_randn(c(1L, lat_ch, h_lat, w_lat), dtype = f32)
+ latents <- flux_pack_latents(latents)$to(device = device)
+
+ img_ids <- flux_prepare_latent_image_ids(h_lat %/% 2L, w_lat %/% 2L)
+ txt_ids <- torch::torch_zeros(txt_len, 3L)
+ ids <- torch::torch_cat(list(txt_ids, img_ids), dim = 1L)
+ rope <- flux_pos_embed(ids,
+ axes_dim = pipeline$transformer$axes_dims_rope %||% c(16L, 56L, 56L))
+ rope <- list(rope[[1]]$to(device = device), rope[[2]]$to(device = device))
+
+ sched <- flowmatch_scheduler_create(
+ shift = pipeline$scheduler_shift %||% 1.0,
+ use_dynamic_shifting = FALSE
+ )
+ n_steps <- as.integer(num_inference_steps)
+ sched <- flowmatch_set_timesteps(
+ sched, n_steps,
+ sigmas = seq(1, 1 / n_steps, length.out = n_steps)
+ )
+
+ # --- Phase 3: denoise --------------------------------------------------------
+ transformer <- onload(pipeline$transformer)
+ if (verbose) {
+ message(sprintf("Denoising: %d steps at %dx%d...", n_steps, width,
+ height))
+ }
+ latents <- .flux_denoise(
+ transformer, latents, sched, prompt_embeds,
+ pooled_prompt_embeds, rope, compute_dtype,
+ chunk_size = pipeline$attn_chunk, verbose = verbose
+ )
+ offload(pipeline$transformer)
+ ltx23_release_dequant_buffers()
+
+ # --- Phase 4: decode -----------------------------------------------------------
+ if (verbose) {
+ message("Decoding...")
+ }
+ latents <- flux_unpack_latents(latents, height, width)
+ 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)
+ }
+ 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)
+
+ img <- img$squeeze(1)$permute(c(2L, 3L, 1L))
+ img <- img$add(1)$div(2)$clamp(0, 1)
+ img_array <- as.array(img)
+
+ gen_seconds <- as.numeric(difftime(Sys.time(), t0, units = "secs"))
+ if (verbose) {
+ message(sprintf("Generated in %.1f s", gen_seconds))
+ }
+
+ if (save_file) {
+ if (is.null(filename)) {
+ filename <- filename_from_prompt(prompt)
+ }
+ save_image(img_array, filename)
+ if (verbose) {
+ message("Saved to ", filename)
+ }
+ }
+
+ metadata <- list(
+ prompt = prompt, width = width, height = height,
+ steps = n_steps, seed = seed, model = "flux1-schnell",
+ precision = pipeline$format, seconds = gen_seconds
+ )
+ invisible(list(image = img_array, metadata = metadata))
+}
diff --git a/R/vae_decoder.R b/R/vae_decoder.R
index bf32dd5..73d76b4 100644
--- a/R/vae_decoder.R
+++ b/R/vae_decoder.R
@@ -8,12 +8,13 @@ VAEResnetBlock <- torch::nn_module(
initialize = function(
in_channels,
- out_channels
+ out_channels,
+ norm_groups = 32
) {
- self$norm1 <- torch::nn_group_norm(32, in_channels, eps = 1e-6)
+ self$norm1 <- torch::nn_group_norm(norm_groups, in_channels, eps = 1e-6)
self$conv1 <- torch::nn_conv2d(in_channels, out_channels,
kernel_size = 3, padding = 1)
- self$norm2 <- torch::nn_group_norm(32, out_channels, eps = 1e-6)
+ self$norm2 <- torch::nn_group_norm(norm_groups, out_channels, eps = 1e-6)
self$conv2 <- torch::nn_conv2d(out_channels, out_channels, kernel_size = 3, padding = 1)
# Shortcut if dimensions change
@@ -49,8 +50,8 @@ VAEResnetBlock <- torch::nn_module(
VAEAttentionBlock <- torch::nn_module(
"VAEAttentionBlock",
- initialize = function(channels) {
- self$group_norm <- torch::nn_group_norm(32, channels, eps = 1e-6)
+ initialize = function(channels, norm_groups = 32) {
+ self$group_norm <- torch::nn_group_norm(norm_groups, channels, eps = 1e-6)
self$to_q <- torch::nn_linear(channels, channels)
self$to_k <- torch::nn_linear(channels, channels)
self$to_v <- torch::nn_linear(channels, channels)
@@ -108,7 +109,8 @@ VAEUpBlock <- torch::nn_module(
in_channels,
out_channels,
num_resnets = 3,
- add_upsample = TRUE
+ add_upsample = TRUE,
+ norm_groups = 32
) {
self$resnets <- torch::nn_module_list()
@@ -118,7 +120,7 @@ VAEUpBlock <- torch::nn_module(
} else {
res_in <- out_channels
}
- self$resnets$append(VAEResnetBlock(res_in, out_channels))
+ self$resnets$append(VAEResnetBlock(res_in, out_channels, norm_groups))
}
if (add_upsample) {
@@ -156,12 +158,13 @@ VAEUpBlock <- torch::nn_module(
VAEMidBlock <- torch::nn_module(
"VAEMidBlock",
- initialize = function(channels) {
+ initialize = function(channels, norm_groups = 32) {
self$resnets <- torch::nn_module_list(list(
- VAEResnetBlock(channels, channels),
- VAEResnetBlock(channels, channels)
+ VAEResnetBlock(channels, channels, norm_groups),
+ VAEResnetBlock(channels, channels, norm_groups)
))
- self$attentions <- torch::nn_module_list(list(VAEAttentionBlock(channels)))
+ self$attentions <- torch::nn_module_list(list(VAEAttentionBlock(channels,
+ norm_groups)))
},
forward = function(x) {
@@ -172,6 +175,62 @@ VAEMidBlock <- torch::nn_module(
}
)
+#' Load HF safetensors VAE weights into the native decoder
+#'
+#' Loads the decoder half of a diffusers AutoencoderKL safetensors file
+#' (e.g. FLUX.1-schnell's \code{vae/diffusion_pytorch_model.safetensors}).
+#' Keys under \code{decoder.} map to the native module 1:1; encoder and
+#' quant-conv keys are skipped (the FLUX VAE has no quant convs, and
+#' txt2img needs no encoder).
+#'
+#' @param native_decoder Native VAE decoder module
+#' @param path Path to the VAE .safetensors file (or a directory
+#' containing diffusion_pytorch_model.safetensors)
+#' @param verbose Print loading progress
+#'
+#' @return The native decoder with loaded weights (invisibly)
+#' @export
+load_decoder_safetensors <- function(native_decoder, path, verbose = TRUE) {
+ path <- path.expand(path)
+ if (dir.exists(path)) {
+ path <- file.path(path, "diffusion_pytorch_model.safetensors")
+ }
+ handle <- safetensors::safetensors$new(path, framework = "torch")
+ keys <- setdiff(handle$keys(), "__metadata__")
+ dec_keys <- keys[startsWith(keys, "decoder.")]
+
+ dests <- native_decoder$named_parameters()
+ filled <- character(0)
+ unmapped <- character(0)
+ torch::with_no_grad({
+ for (key in dec_keys) {
+ native_name <- sub("^decoder\\.", "", key)
+ dest <- dests[[native_name]]
+ if (is.null(dest)) {
+ unmapped <- c(unmapped, key)
+ next
+ }
+ dest$copy_(handle$get_tensor(key))
+ filled <- c(filled, native_name)
+ }
+ })
+
+ unfilled <- setdiff(names(dests), filled)
+ if (length(unmapped)) {
+ stop("VAE decoder load: ", length(unmapped), " unmapped keys, e.g. ",
+ paste(utils::head(unmapped, 3), collapse = ", "))
+ }
+ if (length(unfilled)) {
+ stop("VAE decoder load: ", length(unfilled),
+ " unfilled params, e.g. ",
+ paste(utils::head(unfilled, 3), collapse = ", "))
+ }
+ if (verbose) {
+ message("Loaded ", length(filled), " decoder parameters from ", path)
+ }
+ invisible(native_decoder)
+}
+
#' Load weights from TorchScript decoder into native decoder
#'
#' @param native_decoder Native VAE decoder module
@@ -219,8 +278,13 @@ load_decoder_weights <- function(native_decoder, torchscript_path,
#' Native R torch implementation of the SDXL VAE decoder.
#' Replaces TorchScript decoder for better GPU compatibility.
#'
-#' @param latent_channels Number of latent channels (default 4)
+#' @param latent_channels Number of latent channels (4 for SD/SDXL,
+#' 16 for FLUX/SD3)
#' @param out_channels Number of output channels (default 3 for RGB)
+#' @param block_channels Decoder block channels (reversed encoder
+#' block_out_channels; default matches SD/SDXL and FLUX)
+#' @param norm_groups Group norm groups (default 32; must divide every
+#' entry of \code{block_channels})
#'
#' @return An nn_module representing the VAE decoder
#' @export
@@ -237,37 +301,33 @@ vae_decoder_native <- torch::nn_module(
initialize = function(
latent_channels = 4,
- out_channels = 3
+ out_channels = 3,
+ block_channels = c(512, 512, 256, 128),
+ norm_groups = 32
) {
- # SDXL VAE decoder architecture:
- # Block channels: 512, 512, 256, 128 (reversed from encoder)
- block_channels <- c(512, 512, 256, 128)
+ # Diffusers AutoencoderKL decoder: block channels reversed from the
+ # encoder's block_out_channels; upsamplers on all but the last block.
+ # The SD/SDXL and FLUX/SD3 VAEs share this exact shape (FLUX differs
+ # only in latent_channels = 16).
+ n_blocks <- length(block_channels)
- # Input conv: latent_channels -> 512
- self$conv_in <- torch::nn_conv2d(latent_channels, 512, kernel_size = 3,
- padding = 1)
+ self$conv_in <- torch::nn_conv2d(latent_channels, block_channels[1],
+ kernel_size = 3, padding = 1)
- # Mid block
- self$mid_block <- VAEMidBlock(512)
+ self$mid_block <- VAEMidBlock(block_channels[1], norm_groups)
- # Up blocks (4 blocks, 3 with upsamplers)
self$up_blocks <- torch::nn_module_list()
+ for (i in seq_len(n_blocks)) {
+ in_ch <- block_channels[max(i - 1, 1)]
+ self$up_blocks$append(VAEUpBlock(in_ch, block_channels[i],
+ num_resnets = 3,
+ add_upsample = i < n_blocks,
+ norm_groups = norm_groups))
+ }
- # up_block 0: 512 -> 512, has upsampler
- self$up_blocks$append(VAEUpBlock(512, 512, num_resnets = 3, add_upsample = TRUE))
-
- # up_block 1: 512 -> 512, has upsampler
- self$up_blocks$append(VAEUpBlock(512, 512, num_resnets = 3, add_upsample = TRUE))
-
- # up_block 2: 512 -> 256, has upsampler
- self$up_blocks$append(VAEUpBlock(512, 256, num_resnets = 3, add_upsample = TRUE))
-
- # up_block 3: 256 -> 128, NO upsampler (final block)
- self$up_blocks$append(VAEUpBlock(256, 128, num_resnets = 3, add_upsample = FALSE))
-
- # Output layers
- self$conv_norm_out <- torch::nn_group_norm(32, 128, eps = 1e-6)
- self$conv_out <- torch::nn_conv2d(128, out_channels, kernel_size = 3, padding = 1)
+ last <- block_channels[n_blocks]
+ self$conv_norm_out <- torch::nn_group_norm(norm_groups, last, eps = 1e-6)
+ self$conv_out <- torch::nn_conv2d(last, out_channels, kernel_size = 3, padding = 1)
},
forward = function(x) {
diff --git a/inst/REFERENCES.md b/inst/REFERENCES.md
index f985aba..b0953cd 100644
--- a/inst/REFERENCES.md
+++ b/inst/REFERENCES.md
@@ -1,9 +1,9 @@
-# LTX-2.3 implementation references
+# Implementation references (LTX-2.3, FLUX.1-schnell)
-Every technique in diffuseR's LTX-2.3 support traces to a public,
-permissively licensed source or to our own measured engineering. None of
-it derives from Wan2GP (WanGP Community License) or its `mmgp` module;
-this file documents the actual lineage, idea by idea.
+Every technique in diffuseR's LTX-2.3 and FLUX.1 support traces to a
+public, permissively licensed source or to our own measured engineering.
+None of it derives from Wan2GP (WanGP Community License) or its `mmgp`
+module; this file documents the actual lineage, idea by idea.
## Model architecture and pipeline
@@ -15,6 +15,20 @@ this file documents the actual lineage, idea by idea.
| Single-file checkpoint key layout, embedded config, `model_version` metadata | Format facts of the official Lightricks checkpoints, cross-checked against diffusers `scripts/convert_ltx2_to_diffusers.py` (Apache-2.0) |
| Gemma3 text encoder | Ported from HuggingFace **transformers** (Apache-2.0) |
+## FLUX.1-schnell
+
+| What | Source |
+|---|---|
+| MMDiT transformer (double/single blocks, joint attention, adaLN-Zero variants, RoPE position ids, timestep + pooled-text conditioning) | Ported from HuggingFace **diffusers** (Apache-2.0): `models/transformers/transformer_flux.py`, `models/normalization.py`, `models/embeddings.py` |
+| Pipeline flow: prompt encoding contract, sigma schedule (`linspace(1, 1/N, N)`, static shift), 2x2 latent pack/unpack, latent image ids, VAE scale/shift decode | diffusers `pipelines/flux/pipeline_flux.py` (Apache-2.0) |
+| FlowMatch Euler scheduler | diffusers `schedulers/scheduling_flow_match_euler_discrete.py` (Apache-2.0); shared with the LTX port |
+| 16-channel AutoencoderKL decoder config (no quant convs, scaling 0.3611 / shift 0.1159) | diffusers `scripts/convert_flux_to_diffusers.py` + `convert_sd3_to_diffusers.py` (Apache-2.0) |
+| T5-v1.1 encoder (RMS norms, unscaled unbiased attention, shared relative position bias, gated-GELU FFN) | Ported from HuggingFace **transformers** (Apache-2.0): `models/t5/modeling_t5.py` |
+| CLIP ViT-L text encoder with quick-GELU and argmax-EOS pooling | HuggingFace **transformers** `models/clip/modeling_clip.py` (Apache-2.0); native module shared with the SD/SDXL port |
+| SentencePiece Unigram tokenization (Viterbi best-path over piece log-probs) | Kudo (2018), "Subword Regularization", arXiv:1804.10959; SentencePiece (Apache-2.0); tokenizer.json format facts from HuggingFace tokenizers documentation |
+| NF4/fp8 transformer quantization, cast-set policy, phase offloading, allocator tuning | Same sources as the LTX sections below, applied to the FLUX cast set |
+| Weights | black-forest-labs/FLUX.1-schnell (Apache-2.0; gated HuggingFace repo, downloaded by the user, never redistributed) |
+
## Quantization
| What | Source |
diff --git a/inst/tinytest/fixtures/clip_tiny.safetensors b/inst/tinytest/fixtures/clip_tiny.safetensors
new file mode 100644
index 0000000..7166bc5
Binary files /dev/null and b/inst/tinytest/fixtures/clip_tiny.safetensors differ
diff --git a/inst/tinytest/fixtures/clip_vae_io.safetensors b/inst/tinytest/fixtures/clip_vae_io.safetensors
new file mode 100644
index 0000000..7051b50
Binary files /dev/null and b/inst/tinytest/fixtures/clip_vae_io.safetensors differ
diff --git a/inst/tinytest/fixtures/dit_flux.safetensors b/inst/tinytest/fixtures/dit_flux.safetensors
new file mode 100644
index 0000000..b49a3df
Binary files /dev/null and b/inst/tinytest/fixtures/dit_flux.safetensors differ
diff --git a/inst/tinytest/fixtures/flux_model.safetensors b/inst/tinytest/fixtures/flux_model.safetensors
new file mode 100644
index 0000000..92574c3
Binary files /dev/null and b/inst/tinytest/fixtures/flux_model.safetensors differ
diff --git a/inst/tinytest/fixtures/flux_tiny_ckpt/config.json b/inst/tinytest/fixtures/flux_tiny_ckpt/config.json
new file mode 100644
index 0000000..537acc8
--- /dev/null
+++ b/inst/tinytest/fixtures/flux_tiny_ckpt/config.json
@@ -0,0 +1,19 @@
+{
+ "_class_name": "FluxTransformer2DModel",
+ "_diffusers_version": "0.39.0.dev0",
+ "attention_head_dim": 8,
+ "axes_dims_rope": [
+ 2,
+ 2,
+ 4
+ ],
+ "guidance_embeds": false,
+ "in_channels": 4,
+ "joint_attention_dim": 16,
+ "num_attention_heads": 2,
+ "num_layers": 1,
+ "num_single_layers": 1,
+ "out_channels": null,
+ "patch_size": 1,
+ "pooled_projection_dim": 12
+}
diff --git a/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors
new file mode 100644
index 0000000..fcc4249
Binary files /dev/null and b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00001-of-00003.safetensors differ
diff --git a/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors
new file mode 100644
index 0000000..51bbb41
Binary files /dev/null and b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00002-of-00003.safetensors differ
diff --git a/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors
new file mode 100644
index 0000000..b485a9e
Binary files /dev/null and b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model-00003-of-00003.safetensors differ
diff --git a/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json
new file mode 100644
index 0000000..7ee75de
--- /dev/null
+++ b/inst/tinytest/fixtures/flux_tiny_ckpt/diffusion_pytorch_model.safetensors.index.json
@@ -0,0 +1,69 @@
+{
+ "metadata": {
+ "total_size": 78352
+ },
+ "weight_map": {
+ "context_embedder.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "context_embedder.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "norm_out.linear.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "norm_out.linear.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "proj_out.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "proj_out.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_k.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_q.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_v.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.norm.linear.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.norm.linear.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.proj_mlp.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.proj_mlp.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.proj_out.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "single_transformer_blocks.0.proj_out.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "time_text_embed.text_embedder.linear_1.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.text_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.text_embedder.linear_2.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.text_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.timestep_embedder.linear_1.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.timestep_embedder.linear_1.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.timestep_embedder.linear_2.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "time_text_embed.timestep_embedder.linear_2.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_k_proj.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_k_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_q_proj.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_q_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_v_proj.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.add_v_proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.norm_added_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.norm_added_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.norm_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.norm_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_add_out.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_add_out.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_k.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_k.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_out.0.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_out.0.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_q.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_q.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_v.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.attn.to_v.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff.net.0.proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff.net.2.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff.net.2.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff_context.net.0.proj.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff_context.net.0.proj.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.ff_context.net.2.bias": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "transformer_blocks.0.ff_context.net.2.weight": "diffusion_pytorch_model-00003-of-00003.safetensors",
+ "transformer_blocks.0.norm1.linear.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "transformer_blocks.0.norm1.linear.weight": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "transformer_blocks.0.norm1_context.linear.bias": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "transformer_blocks.0.norm1_context.linear.weight": "diffusion_pytorch_model-00002-of-00003.safetensors",
+ "x_embedder.bias": "diffusion_pytorch_model-00001-of-00003.safetensors",
+ "x_embedder.weight": "diffusion_pytorch_model-00001-of-00003.safetensors"
+ }
+}
diff --git a/inst/tinytest/fixtures/rope_flux.safetensors b/inst/tinytest/fixtures/rope_flux.safetensors
new file mode 100644
index 0000000..5324754
Binary files /dev/null and b/inst/tinytest/fixtures/rope_flux.safetensors differ
diff --git a/inst/tinytest/fixtures/t5_flux.safetensors b/inst/tinytest/fixtures/t5_flux.safetensors
new file mode 100644
index 0000000..d89dba0
Binary files /dev/null and b/inst/tinytest/fixtures/t5_flux.safetensors differ
diff --git a/inst/tinytest/fixtures/t5_tiny_ckpt/config.json b/inst/tinytest/fixtures/t5_tiny_ckpt/config.json
new file mode 100644
index 0000000..b03b777
--- /dev/null
+++ b/inst/tinytest/fixtures/t5_tiny_ckpt/config.json
@@ -0,0 +1,31 @@
+{
+ "architectures": [
+ "T5EncoderModel"
+ ],
+ "classifier_dropout": 0.0,
+ "d_ff": 32,
+ "d_kv": 4,
+ "d_model": 16,
+ "dense_act_fn": "gelu_new",
+ "dropout_rate": 0.0,
+ "dtype": "float32",
+ "eos_token_id": 1,
+ "feed_forward_proj": "gated-gelu",
+ "initializer_factor": 1.0,
+ "is_decoder": false,
+ "is_encoder_decoder": false,
+ "is_gated_act": true,
+ "layer_norm_epsilon": 1e-06,
+ "model_type": "t5",
+ "num_decoder_layers": 2,
+ "num_heads": 4,
+ "num_layers": 2,
+ "pad_token_id": 0,
+ "relative_attention_max_distance": 128,
+ "relative_attention_num_buckets": 32,
+ "scale_decoder_outputs": true,
+ "tie_word_embeddings": true,
+ "transformers_version": "5.13.0",
+ "use_cache": false,
+ "vocab_size": 100
+}
diff --git a/inst/tinytest/fixtures/t5_tiny_ckpt/model.safetensors b/inst/tinytest/fixtures/t5_tiny_ckpt/model.safetensors
new file mode 100644
index 0000000..313af14
Binary files /dev/null and b/inst/tinytest/fixtures/t5_tiny_ckpt/model.safetensors differ
diff --git a/inst/tinytest/fixtures/t5_tokenizer_cases.json b/inst/tinytest/fixtures/t5_tokenizer_cases.json
new file mode 100644
index 0000000..e9bfe19
--- /dev/null
+++ b/inst/tinytest/fixtures/t5_tokenizer_cases.json
@@ -0,0 +1,731 @@
+{
+ "cases": [
+ {
+ "text": "a photo of a cat",
+ "ids": [
+ 3,
+ 9,
+ 1202,
+ 13,
+ 3,
+ 9,
+ 1712,
+ 1
+ ]
+ },
+ {
+ "text": "A sunset over mountains, ultra detailed, 8k",
+ "ids": [
+ 71,
+ 13744,
+ 147,
+ 8022,
+ 6,
+ 6173,
+ 3117,
+ 6,
+ 505,
+ 157,
+ 1
+ ]
+ },
+ {
+ "text": "Hello, world!",
+ "ids": [
+ 8774,
+ 6,
+ 296,
+ 55,
+ 1
+ ]
+ },
+ {
+ "text": "The quick brown fox jumps over the lazy dog.",
+ "ids": [
+ 37,
+ 1704,
+ 4216,
+ 3,
+ 20400,
+ 4418,
+ 7,
+ 147,
+ 8,
+ 19743,
+ 1782,
+ 5,
+ 1
+ ]
+ },
+ {
+ "text": "it's a beautiful day; isn't it?",
+ "ids": [
+ 34,
+ 31,
+ 7,
+ 3,
+ 9,
+ 786,
+ 239,
+ 117,
+ 19,
+ 29,
+ 31,
+ 17,
+ 34,
+ 58,
+ 1
+ ]
+ },
+ {
+ "text": "3.14159 and 2,000,000 dollars",
+ "ids": [
+ 1877,
+ 2534,
+ 27904,
+ 11,
+ 204,
+ 23916,
+ 3740,
+ 1
+ ]
+ },
+ {
+ "text": "state-of-the-art text-to-image generation",
+ "ids": [
+ 538,
+ 18,
+ 858,
+ 18,
+ 532,
+ 18,
+ 1408,
+ 1499,
+ 18,
+ 235,
+ 18,
+ 8221,
+ 3381,
+ 1
+ ]
+ },
+ {
+ "text": " leading spaces",
+ "ids": [
+ 1374,
+ 4856,
+ 1
+ ]
+ },
+ {
+ "text": "trailing spaces ",
+ "ids": [
+ 5032,
+ 53,
+ 4856,
+ 1
+ ]
+ },
+ {
+ "text": "double and triple spaces",
+ "ids": [
+ 1486,
+ 11,
+ 12063,
+ 4856,
+ 1
+ ]
+ },
+ {
+ "text": "UPPERCASE lowercase MiXeD",
+ "ids": [
+ 3,
+ 6880,
+ 8742,
+ 254,
+ 17892,
+ 1364,
+ 6701,
+ 2133,
+ 4,
+ 15,
+ 308,
+ 1
+ ]
+ },
+ {
+ "text": "email@example.com and https://example.org/path?q=1",
+ "ids": [
+ 791,
+ 1741,
+ 994,
+ 9,
+ 9208,
+ 5,
+ 287,
+ 11,
+ 4893,
+ 1303,
+ 994,
+ 9,
+ 9208,
+ 5,
+ 1677,
+ 87,
+ 8292,
+ 58,
+ 1824,
+ 2423,
+ 536,
+ 1
+ ]
+ },
+ {
+ "text": "quotes \"double\" and 'single'",
+ "ids": [
+ 7599,
+ 96,
+ 25761,
+ 121,
+ 11,
+ 3,
+ 31,
+ 7,
+ 53,
+ 109,
+ 31,
+ 1
+ ]
+ },
+ {
+ "text": "(parentheses) [brackets] {braces}",
+ "ids": [
+ 41,
+ 12352,
+ 88,
+ 2260,
+ 61,
+ 784,
+ 1939,
+ 8849,
+ 17,
+ 7,
+ 908,
+ 3,
+ 2,
+ 1939,
+ 2319,
+ 2,
+ 1
+ ]
+ },
+ {
+ "text": "semi;colon co:lon sla/sh back\\slash",
+ "ids": [
+ 4772,
+ 117,
+ 8135,
+ 29,
+ 576,
+ 10,
+ 40,
+ 106,
+ 3,
+ 7,
+ 521,
+ 87,
+ 7,
+ 107,
+ 223,
+ 2,
+ 7,
+ 8058,
+ 1
+ ]
+ },
+ {
+ "text": "underscores_and_snake_case",
+ "ids": [
+ 31530,
+ 7,
+ 834,
+ 232,
+ 834,
+ 7,
+ 29,
+ 9,
+ 1050,
+ 834,
+ 6701,
+ 1
+ ]
+ },
+ {
+ "text": "a",
+ "ids": [
+ 3,
+ 9,
+ 1
+ ]
+ },
+ {
+ "text": "",
+ "ids": [
+ 1
+ ]
+ },
+ {
+ "text": "café résumé naïve",
+ "ids": [
+ 11949,
+ 1417,
+ 4078,
+ 154,
+ 3,
+ 29,
+ 9,
+ 2,
+ 162,
+ 1
+ ]
+ },
+ {
+ "text": "em—dash and en–dash",
+ "ids": [
+ 3,
+ 15,
+ 51,
+ 318,
+ 26,
+ 3198,
+ 11,
+ 3,
+ 35,
+ 104,
+ 26,
+ 3198,
+ 1
+ ]
+ },
+ {
+ "text": "100% of $50 + €20",
+ "ids": [
+ 2349,
+ 13,
+ 13309,
+ 1768,
+ 3416,
+ 1755,
+ 1
+ ]
+ },
+ {
+ "text": "An astronaut riding a horse on Mars, photorealistic",
+ "ids": [
+ 389,
+ 30059,
+ 7494,
+ 3,
+ 9,
+ 4952,
+ 30,
+ 11856,
+ 6,
+ 1202,
+ 6644,
+ 3040,
+ 1
+ ]
+ },
+ {
+ "text": "watercolor painting of a fox in a snowy forest",
+ "ids": [
+ 27610,
+ 3924,
+ 13,
+ 3,
+ 9,
+ 3,
+ 20400,
+ 16,
+ 3,
+ 9,
+ 2983,
+ 63,
+ 5827,
+ 1
+ ]
+ },
+ {
+ "text": "the mitochondria is the powerhouse of the cell",
+ "ids": [
+ 8,
+ 28642,
+ 19,
+ 8,
+ 579,
+ 1840,
+ 13,
+ 8,
+ 2358,
+ 1
+ ]
+ },
+ {
+ "text": "The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. ",
+ "ids": [
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 16,
+ 5932,
+ 7,
+ 5,
+ 1
+ ]
+ }
+ ],
+ "padded": [
+ {
+ "text": "a photo of a cat",
+ "max_length": 16,
+ "ids": [
+ 3,
+ 9,
+ 1202,
+ 13,
+ 3,
+ 9,
+ 1712,
+ 1,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0
+ ],
+ "mask": [
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0,
+ 0
+ ]
+ },
+ {
+ "text": "The quick brown fox jumps over the lazy dog.",
+ "max_length": 16,
+ "ids": [
+ 37,
+ 1704,
+ 4216,
+ 3,
+ 20400,
+ 4418,
+ 7,
+ 147,
+ 8,
+ 19743,
+ 1782,
+ 5,
+ 1,
+ 0,
+ 0,
+ 0
+ ],
+ "mask": [
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 0,
+ 0,
+ 0
+ ]
+ },
+ {
+ "text": "The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. The transformer architecture uses self-attention mechanisms to model long-range dependencies in sequences. ",
+ "max_length": 16,
+ "ids": [
+ 37,
+ 19903,
+ 4648,
+ 2284,
+ 1044,
+ 18,
+ 25615,
+ 12009,
+ 12,
+ 825,
+ 307,
+ 18,
+ 5517,
+ 6002,
+ 11573,
+ 1
+ ],
+ "mask": [
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1,
+ 1
+ ]
+ }
+ ]
+}
\ No newline at end of file
diff --git a/inst/tinytest/fixtures/vae16_tiny.safetensors b/inst/tinytest/fixtures/vae16_tiny.safetensors
new file mode 100644
index 0000000..6c04633
Binary files /dev/null and b/inst/tinytest/fixtures/vae16_tiny.safetensors differ
diff --git a/inst/tinytest/test_clip_vae_flux.R b/inst/tinytest/test_clip_vae_flux.R
new file mode 100644
index 0000000..4c3b0ea
--- /dev/null
+++ b/inst/tinytest/test_clip_vae_flux.R
@@ -0,0 +1,62 @@
+# Parity tests for the FLUX CLIP extensions (quick_gelu, safetensors
+# loader, pooled output) and the 16-channel VAE decoder (fixtures from
+# tools/gen_fixtures_flux_clip_vae.py, checked in).
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+fixture <- function(name) {
+ p <- system.file("tinytest", "fixtures", name, package = "diffuseR")
+ if (p == "") p <- file.path("fixtures", name)
+ p
+}
+io_path <- fixture("clip_vae_io.safetensors")
+clip_path <- fixture("clip_tiny.safetensors")
+vae_path <- fixture("vae16_tiny.safetensors")
+if (!file.exists(io_path)) exit_file("clip/vae fixtures missing")
+
+fx <- safetensors::safe_load_file(io_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+
+# --- CLIP: quick_gelu forward + argmax pooled output -------------------------------
+
+enc <- text_encoder_native(
+ vocab_size = 1000L, context_length = 77L, embed_dim = 16L,
+ num_layers = 2L, num_heads = 2L, mlp_dim = 32L,
+ apply_final_ln = TRUE, gelu_type = "quick"
+)
+enc$eval()
+load_text_encoder_safetensors(enc, clip_path, verbose = FALSE)
+
+ids <- fx$clip_input_ids$to(dtype = torch::torch_long())
+hidden <- torch::with_no_grad(enc(ids))
+expect_true(max_abs_diff(hidden, fx$clip_last_hidden) < 1e-5)
+
+pooled <- clip_pooled_output(hidden, ids)
+expect_equal(as.integer(pooled$shape), c(2L, 16L))
+expect_true(max_abs_diff(pooled, fx$clip_pooled) < 1e-5)
+
+# --- 16-channel VAE decoder ----------------------------------------------------------
+
+dec <- vae_decoder_native(
+ latent_channels = 16L,
+ block_channels = c(32L, 32L, 16L, 8L),
+ norm_groups = 8L
+)
+dec$eval()
+load_decoder_safetensors(dec, vae_path, verbose = FALSE)
+
+img <- torch::with_no_grad(dec(fx$vae_latent))
+expect_equal(as.integer(img$shape), as.integer(fx$vae_image$shape))
+expect_true(max_abs_diff(img, fx$vae_image) < 1e-5)
diff --git a/inst/tinytest/test_dit_flux.R b/inst/tinytest/test_dit_flux.R
new file mode 100644
index 0000000..9c36d51
--- /dev/null
+++ b/inst/tinytest/test_dit_flux.R
@@ -0,0 +1,143 @@
+# Parity tests for the FLUX transformer blocks against diffusers
+# reference fixtures (generated by tools/gen_fixtures_flux_dit.py,
+# checked in).
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+fixture_path <- system.file("tinytest", "fixtures", "dit_flux.safetensors",
+ package = "diffuseR")
+if (fixture_path == "") fixture_path <- "fixtures/dit_flux.safetensors"
+if (!file.exists(fixture_path)) exit_file("flux dit fixtures missing")
+
+fx <- safetensors::safe_load_file(fixture_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+
+# Copy fixture weights (diffusers state-dict names) into an R module.
+# Errors if any name has no destination or any destination stays unfilled.
+load_named_weights <- function(module, weights) {
+ dests <- c(module$named_parameters(), module$named_buffers())
+ missing_dest <- setdiff(names(weights), names(dests))
+ if (length(missing_dest)) {
+ stop("No destination for: ", paste(utils::head(missing_dest, 5), collapse = ", "))
+ }
+ unfilled <- setdiff(names(dests), names(weights))
+ if (length(unfilled)) {
+ stop("Unfilled params: ", paste(utils::head(unfilled, 5), collapse = ", "))
+ }
+ torch::with_no_grad({
+ for (name in names(weights)) dests[[name]]$copy_(weights[[name]])
+ })
+ invisible(module)
+}
+
+fixture_group <- function(prefix) {
+ keys <- grep(paste0("^", prefix, "\\."), names(fx), value = TRUE)
+ w <- fx[keys]
+ names(w) <- sub(paste0("^", prefix, "\\."), "", keys)
+ w
+}
+
+rope <- list(fx$rope_cos, fx$rope_sin)
+DIM <- 16L
+HEADS <- 2L
+HEAD_DIM <- 8L
+S_TXT <- 7L
+
+# --- flux_ada_layer_norm_zero ---------------------------------------------------
+
+adazero <- flux_ada_layer_norm_zero(DIM)
+adazero$eval()
+load_named_weights(adazero, fixture_group("adazero"))
+az <- torch::with_no_grad(adazero(fx$joint_x, emb = fx$temb))
+expect_true(max_abs_diff(az[[1]], fx$adazero_x_norm) < 1e-5)
+expect_true(max_abs_diff(az[[2]], fx$adazero_gate_msa) < 1e-5)
+expect_true(max_abs_diff(az[[3]], fx$adazero_shift_mlp) < 1e-5)
+expect_true(max_abs_diff(az[[4]], fx$adazero_scale_mlp) < 1e-5)
+expect_true(max_abs_diff(az[[5]], fx$adazero_gate_mlp) < 1e-5)
+
+# --- flux_ada_layer_norm_zero_single ----------------------------------------------
+
+adasingle <- flux_ada_layer_norm_zero_single(DIM)
+adasingle$eval()
+load_named_weights(adasingle, fixture_group("adasingle"))
+as_ <- torch::with_no_grad(adasingle(fx$joint_x, emb = fx$temb))
+expect_true(max_abs_diff(as_[[1]], fx$adasingle_x_norm) < 1e-5)
+expect_true(max_abs_diff(as_[[2]], fx$adasingle_gate) < 1e-5)
+
+# --- flux_ada_layer_norm_continuous (scale-first chunk order) ----------------------
+
+adacont <- flux_ada_layer_norm_continuous(DIM, DIM)
+adacont$eval()
+load_named_weights(adacont, fixture_group("adacont"))
+ac <- torch::with_no_grad(adacont(fx$joint_x, fx$temb))
+expect_true(max_abs_diff(ac, fx$adacont_out) < 1e-5)
+
+# --- flux_attention, double-stream variant ------------------------------------------
+
+attn_d <- flux_attention(DIM, HEADS, HEAD_DIM, added_kv = TRUE)
+attn_d$eval()
+load_named_weights(attn_d, fixture_group("attnd"))
+ad <- torch::with_no_grad(attn_d(
+ hidden_states = fx$img_x,
+ encoder_hidden_states = fx$txt_x,
+ image_rotary_emb = rope
+))
+expect_equal(as.integer(ad[[1]]$shape), as.integer(fx$attnd_out$shape))
+expect_true(max_abs_diff(ad[[1]], fx$attnd_out) < 1e-5)
+expect_true(max_abs_diff(ad[[2]], fx$attnd_ctx_out) < 1e-5)
+
+# --- flux_attention, pre_only variant -------------------------------------------------
+
+attn_s <- flux_attention(DIM, HEADS, HEAD_DIM, pre_only = TRUE)
+attn_s$eval()
+load_named_weights(attn_s, fixture_group("attns"))
+asn <- torch::with_no_grad(attn_s(
+ hidden_states = fx$joint_x,
+ image_rotary_emb = rope
+))
+expect_true(max_abs_diff(asn, fx$attns_out) < 1e-5)
+
+# --- flux_double_block -----------------------------------------------------------------
+
+dbl <- flux_double_block(DIM, HEADS, HEAD_DIM)
+dbl$eval()
+load_named_weights(dbl, fixture_group("dbl"))
+db <- torch::with_no_grad(dbl(
+ hidden_states = fx$img_x,
+ encoder_hidden_states = fx$txt_x,
+ temb = fx$temb,
+ image_rotary_emb = rope
+))
+expect_true(max_abs_diff(db[[1]], fx$dbl_enc_out) < 1e-5)
+expect_true(max_abs_diff(db[[2]], fx$dbl_hid_out) < 1e-5)
+
+# --- flux_single_block -------------------------------------------------------------------
+# The R block takes the pre-concatenated [txt; img] sequence and returns it
+# joint; the reference returns the split streams.
+
+sgl <- flux_single_block(DIM, HEADS, HEAD_DIM)
+sgl$eval()
+load_named_weights(sgl, fixture_group("sgl"))
+joint_in <- torch::torch_cat(list(fx$txt_x, fx$img_x), dim = 2L)
+sg <- torch::with_no_grad(sgl(
+ hidden_states = joint_in,
+ temb = fx$temb,
+ image_rotary_emb = rope
+))
+expect_true(max_abs_diff(sg$narrow(2L, 1L, S_TXT), fx$sgl_enc_out) < 1e-5)
+expect_true(max_abs_diff(
+ sg$narrow(2L, S_TXT + 1L, sg$shape[2] - S_TXT),
+ fx$sgl_hid_out
+) < 1e-5)
diff --git a/inst/tinytest/test_flux_transformer.R b/inst/tinytest/test_flux_transformer.R
new file mode 100644
index 0000000..9784e5a
--- /dev/null
+++ b/inst/tinytest/test_flux_transformer.R
@@ -0,0 +1,61 @@
+# Parity test for the full FLUX transformer against a tiny random-init
+# diffusers FluxTransformer2DModel (fixture generated by
+# tools/gen_fixtures_flux_model.py, checked in).
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+fixture_path <- system.file("tinytest", "fixtures", "flux_model.safetensors",
+ package = "diffuseR")
+if (fixture_path == "") fixture_path <- "fixtures/flux_model.safetensors"
+if (!file.exists(fixture_path)) exit_file("flux model fixtures missing")
+
+fx <- safetensors::safe_load_file(fixture_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+
+model <- flux_transformer(
+ in_channels = 4L,
+ num_layers = 1L,
+ num_single_layers = 1L,
+ attention_head_dim = 8L,
+ num_attention_heads = 2L,
+ joint_attention_dim = 16L,
+ pooled_projection_dim = 12L,
+ axes_dims_rope = c(2L, 2L, 4L)
+)
+model$eval()
+
+# Strict both ways: every fixture key must land, every param must fill.
+# This proves the R module tree matches the diffusers state dict exactly.
+weights <- fx[grep("^model\\.", names(fx))]
+names(weights) <- sub("^model\\.", "", names(weights))
+dests <- c(model$named_parameters(), model$named_buffers())
+expect_equal(character(0), setdiff(names(weights), names(dests)))
+expect_equal(character(0), setdiff(names(dests), names(weights)))
+torch::with_no_grad({
+ for (name in names(weights)) dests[[name]]$copy_(weights[[name]])
+})
+
+ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L)
+rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 4L))
+
+out <- torch::with_no_grad(model(
+ hidden_states = fx$hidden,
+ encoder_hidden_states = fx$encoder,
+ pooled_projections = fx$pooled,
+ timestep = fx$timestep,
+ image_rotary_emb = rope
+))
+expect_equal(as.integer(out$shape), as.integer(fx$out$shape))
+expect_true(max_abs_diff(out, fx$out) < 1e-4)
diff --git a/inst/tinytest/test_quantize_flux.R b/inst/tinytest/test_quantize_flux.R
new file mode 100644
index 0000000..03b8e02
--- /dev/null
+++ b/inst/tinytest/test_quantize_flux.R
@@ -0,0 +1,119 @@
+# FLUX checkpoint opening + quantization round trip on the tiny sharded
+# diffusers checkpoint (fixture generated by tools/gen_fixtures_flux_model.py,
+# checked in). Everything runs on CPU.
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+ckpt_dir <- system.file("tinytest", "fixtures", "flux_tiny_ckpt",
+ package = "diffuseR")
+if (ckpt_dir == "") ckpt_dir <- "fixtures/flux_tiny_ckpt"
+if (!dir.exists(ckpt_dir)) exit_file("flux tiny checkpoint missing")
+
+fixture_path <- system.file("tinytest", "fixtures", "flux_model.safetensors",
+ package = "diffuseR")
+if (fixture_path == "") fixture_path <- "fixtures/flux_model.safetensors"
+if (!file.exists(fixture_path)) exit_file("flux model fixtures missing")
+
+fx <- safetensors::safe_load_file(fixture_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+cosine_sim <- function(a, b) {
+ a <- a$to(dtype = torch::torch_float32())$flatten()
+ b <- b$to(dtype = torch::torch_float32())$flatten()
+ as.numeric(torch::torch_dot(a, b) / (a$norm() * b$norm()))
+}
+
+# --- cast-set census ------------------------------------------------------------
+# Tiny config: 1 double block x 14 + 1 single block x 6 = 20 cast weights.
+# (Full schnell: 19 x 14 + 38 x 6 = 494.)
+
+ckpt <- flux_open_checkpoint(ckpt_dir)
+expect_true(inherits(ckpt, "ltx23_checkpoint"))
+expect_equal(sum(flux_is_quant_key(ckpt$keys)), 20L)
+expect_equal(ckpt$config$num_layers, 1L)
+
+# The sharded index resolves every key to a readable tensor
+t0 <- ckpt$handle$get_tensor("x_embedder.weight")
+expect_equal(as.integer(t0$shape), c(16L, 4L))
+
+# --- full-precision load through the sharded index --------------------------------
+
+model <- flux_load_transformer(ckpt, device = "cpu", dtype = "float32",
+ verbose = FALSE)
+ids <- torch::torch_cat(list(fx$txt_ids, fx$img_ids), dim = 1L)
+rope <- flux_pos_embed(ids, axes_dim = c(2L, 2L, 4L))
+out_full <- torch::with_no_grad(model(
+ hidden_states = fx$hidden,
+ encoder_hidden_states = fx$encoder,
+ pooled_projections = fx$pooled,
+ timestep = fx$timestep,
+ image_rotary_emb = rope
+))
+expect_true(max_abs_diff(out_full, fx$out) < 1e-4)
+
+# --- NF4 round trip -----------------------------------------------------------------
+
+nf4_dir <- file.path(tempdir(), "flux-tiny-nf4")
+unlink(nf4_dir, recursive = TRUE)
+manifest <- flux_quantize(ckpt_dir, output_dir = nf4_dir, format = "nf4",
+ verbose = FALSE)
+expect_equal(manifest$cast, 20L)
+expect_true(file.exists(file.path(nf4_dir, "manifest.json")))
+
+nf4_ckpt <- flux_open_quantized(nf4_dir)
+expect_equal(nf4_ckpt$format, "nf4")
+model_nf4 <- flux_load_transformer(nf4_ckpt, device = "cpu", verbose = FALSE)
+out_nf4 <- torch::with_no_grad(model_nf4(
+ hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()),
+ encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()),
+ pooled_projections = fx$pooled$to(dtype = torch::torch_bfloat16()),
+ timestep = fx$timestep,
+ image_rotary_emb = rope
+))
+expect_true(all(is.finite(as.numeric(out_nf4$to(dtype = torch::torch_float32())))))
+expect_true(cosine_sim(out_nf4, out_full) > 0.98)
+ltx23_release_dequant_buffers()
+
+# --- fp8 round trip (needs an F8-capable safetensors build) ---------------------------
+
+f8_ok <- tryCatch({
+ x <- torch::torch_randn(2, 2)$to(dtype = torch::torch_float8_e4m3fn())
+ tmp <- tempfile(fileext = ".safetensors")
+ safetensors::safe_save_file(list(w = x), tmp)
+ y <- safetensors::safe_load_file(tmp, framework = "torch")
+ unlink(tmp)
+ TRUE
+}, error = function(e) FALSE)
+
+if (f8_ok) {
+ fp8_dir <- file.path(tempdir(), "flux-tiny-fp8")
+ unlink(fp8_dir, recursive = TRUE)
+ manifest8 <- flux_quantize(ckpt_dir, output_dir = fp8_dir, format = "fp8",
+ verbose = FALSE)
+ expect_equal(manifest8$cast, 20L)
+
+ fp8_ckpt <- flux_open_quantized(fp8_dir)
+ model_fp8 <- flux_load_transformer(fp8_ckpt, device = "cpu", pin = FALSE,
+ verbose = FALSE)
+ out_fp8 <- torch::with_no_grad(model_fp8(
+ hidden_states = fx$hidden$to(dtype = torch::torch_bfloat16()),
+ encoder_hidden_states = fx$encoder$to(dtype = torch::torch_bfloat16()),
+ pooled_projections = fx$pooled$to(dtype = torch::torch_bfloat16()),
+ timestep = fx$timestep,
+ image_rotary_emb = rope
+ ))
+ expect_true(cosine_sim(out_fp8, out_full) > 0.99)
+}
+
+options(diffuseR.block_gc = NULL)
diff --git a/inst/tinytest/test_rope_flux.R b/inst/tinytest/test_rope_flux.R
new file mode 100644
index 0000000..be96d0d
--- /dev/null
+++ b/inst/tinytest/test_rope_flux.R
@@ -0,0 +1,79 @@
+# Parity tests for the FLUX RoPE port and latent pack/unpack against
+# diffusers reference fixtures (generated by tools/gen_fixtures_flux.py,
+# checked in).
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+fixture_path <- system.file("tinytest", "fixtures", "rope_flux.safetensors",
+ package = "diffuseR")
+if (fixture_path == "") {
+ # Running from the source tree (e.g. during development)
+ fixture_path <- "fixtures/rope_flux.safetensors"
+}
+if (!file.exists(fixture_path)) exit_file("flux rope fixtures missing")
+
+fx <- safetensors::safe_load_file(fixture_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+
+# --- latent image ids ---------------------------------------------------------
+
+ids_r <- flux_prepare_latent_image_ids(8L, 12L)
+expect_equal(as.integer(ids_r$shape), as.integer(fx$img_ids$shape))
+expect_true(max_abs_diff(ids_r, fx$img_ids) == 0)
+
+# --- flux_pos_embed: full-size axes (16, 56, 56) -------------------------------
+
+pf <- flux_pos_embed(fx$ids, axes_dim = c(16L, 56L, 56L))
+expect_equal(as.integer(pf[[1]]$shape), as.integer(fx$pos_full_cos$shape))
+expect_true(max_abs_diff(pf[[1]], fx$pos_full_cos) < 1e-6)
+expect_true(max_abs_diff(pf[[2]], fx$pos_full_sin) < 1e-6)
+
+# Text ids (all zero) get the identity rotation: cos 1, sin 0
+txt_cos <- pf[[1]][1:7, ]
+txt_sin <- pf[[2]][1:7, ]
+expect_true(max_abs_diff(txt_cos, torch::torch_ones_like(txt_cos)) == 0)
+expect_true(max_abs_diff(txt_sin, torch::torch_zeros_like(txt_sin)) == 0)
+
+# --- flux_pos_embed: tiny axes (2, 2, 4) ---------------------------------------
+
+pt <- flux_pos_embed(fx$ids, axes_dim = c(2L, 2L, 4L))
+expect_equal(as.integer(pt[[1]]$shape), as.integer(fx$pos_tiny_cos$shape))
+expect_true(max_abs_diff(pt[[1]], fx$pos_tiny_cos) < 1e-6)
+expect_true(max_abs_diff(pt[[2]], fx$pos_tiny_sin) < 1e-6)
+
+# --- flux_apply_rotary_emb ------------------------------------------------------
+
+rot <- flux_apply_rotary_emb(fx$rot_x, list(fx$pos_tiny_cos, fx$pos_tiny_sin))
+expect_equal(as.integer(rot$shape), as.integer(fx$rot_out$shape))
+expect_true(max_abs_diff(rot, fx$rot_out) < 1e-5)
+
+# fp16 input stays fp16 and matches reference values
+rot16 <- flux_apply_rotary_emb(fx$rot_x_f16, list(fx$pos_tiny_cos, fx$pos_tiny_sin))
+expect_equal(rot16$dtype$.type(), "Half")
+expect_true(max_abs_diff(rot16, fx$rot_out_f16) < 1e-3)
+
+# --- latent pack / unpack -------------------------------------------------------
+
+packed <- flux_pack_latents(fx$pack_in)
+expect_equal(as.integer(packed$shape), as.integer(fx$pack_out$shape))
+expect_true(max_abs_diff(packed, fx$pack_out) == 0)
+
+unpacked <- flux_unpack_latents(fx$pack_out, height = 64L, width = 96L,
+ vae_scale_factor = 8L)
+expect_equal(as.integer(unpacked$shape), as.integer(fx$unpack_out$shape))
+expect_true(max_abs_diff(unpacked, fx$unpack_out) == 0)
+
+# R-side roundtrip identity: unpack(pack(x)) == x
+expect_true(max_abs_diff(unpacked, fx$pack_in) == 0)
diff --git a/inst/tinytest/test_t5_flux.R b/inst/tinytest/test_t5_flux.R
new file mode 100644
index 0000000..b2fca59
--- /dev/null
+++ b/inst/tinytest/test_t5_flux.R
@@ -0,0 +1,52 @@
+# Parity tests for the T5 encoder port against HF transformers reference
+# fixtures (generated by tools/gen_fixtures_t5.py, checked in).
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+if (!requireNamespace("safetensors", quietly = TRUE)) {
+ exit_file("safetensors not installed")
+}
+
+library(diffuseR)
+
+fixture_path <- system.file("tinytest", "fixtures", "t5_flux.safetensors",
+ package = "diffuseR")
+if (fixture_path == "") fixture_path <- "fixtures/t5_flux.safetensors"
+if (!file.exists(fixture_path)) exit_file("t5 fixtures missing")
+
+ckpt_dir <- system.file("tinytest", "fixtures", "t5_tiny_ckpt",
+ package = "diffuseR")
+if (ckpt_dir == "") ckpt_dir <- "fixtures/t5_tiny_ckpt"
+if (!dir.exists(ckpt_dir)) exit_file("t5 tiny checkpoint missing")
+
+fx <- safetensors::safe_load_file(fixture_path, framework = "torch")
+
+max_abs_diff <- function(a, b) {
+ as.numeric(torch::torch_max(torch::torch_abs(
+ a$to(dtype = torch::torch_float32()) - b$to(dtype = torch::torch_float32())
+ )))
+}
+
+# --- relative position bucketing: exact integer parity ---------------------------
+
+buckets_r <- diffuseR:::.t5_relative_position_bucket(
+ fx$rel_positions$to(dtype = torch::torch_long())
+)
+expect_equal(
+ as.integer(buckets_r$flatten()),
+ as.integer(fx$rel_buckets$flatten())
+)
+
+# --- tiny encoder parity (padded batch, no mask) -----------------------------------
+
+model <- load_t5_text_encoder(ckpt_dir, device = "cpu", dtype = "float32",
+ verbose = FALSE)
+out <- torch::with_no_grad(
+ model(fx$input_ids$to(dtype = torch::torch_long()) + 1L)
+)
+expect_equal(as.integer(out$shape), as.integer(fx$out$shape))
+expect_true(max_abs_diff(out, fx$out) < 1e-4)
+
+# (encode_with_t5 is exercised end-to-end in the pipeline tests; real
+# token ids exceed this fixture's 100-entry vocab.)
diff --git a/inst/tinytest/test_tokenizer_unigram.R b/inst/tinytest/test_tokenizer_unigram.R
new file mode 100644
index 0000000..c3ea87b
--- /dev/null
+++ b/inst/tinytest/test_tokenizer_unigram.R
@@ -0,0 +1,78 @@
+# Exact token-id parity for the Unigram (T5/SentencePiece) tokenizer
+# against HuggingFace tokenizers reference cases (generated by
+# tools/gen_t5_tokenizer_cases.py, checked in).
+#
+# The tokenizer.json itself is not shipped (2.4 MB, and FLUX's copy is
+# behind the gated repo); the test locates one and skips if absent.
+
+library(diffuseR)
+
+find_t5_tokenizer <- function() {
+ # 1. Explicit override
+ p <- Sys.getenv("DIFFUSER_T5_TOKENIZER", "")
+ if (nzchar(p) && file.exists(p)) {
+ return(p)
+ }
+ # 2. FLUX.1-schnell HF cache (post-download_flux1)
+ if (requireNamespace("hfhub", quietly = TRUE)) {
+ p <- tryCatch(
+ suppressMessages(hfhub::hub_download(
+ "black-forest-labs/FLUX.1-schnell",
+ "tokenizer_2/tokenizer.json", local_files_only = TRUE
+ )),
+ error = function(e) ""
+ )
+ if (nzchar(p) && file.exists(p)) {
+ return(p)
+ }
+ }
+ # 3. Dev copy in the source tree (tools/gen_t5_tokenizer_cases.py)
+ p <- "../../tools/cache/tokenizer_t5.json"
+ if (file.exists(p)) {
+ return(p)
+ }
+ ""
+}
+
+tok_path <- find_t5_tokenizer()
+if (!nzchar(tok_path)) exit_file("no T5 tokenizer.json available")
+
+cases_path <- system.file("tinytest", "fixtures", "t5_tokenizer_cases.json",
+ package = "diffuseR")
+if (cases_path == "") cases_path <- "fixtures/t5_tokenizer_cases.json"
+if (!file.exists(cases_path)) exit_file("t5 tokenizer cases missing")
+
+cases <- jsonlite::fromJSON(cases_path, simplifyVector = FALSE)
+
+tok <- unigram_tokenizer(tok_path)
+expect_equal(tok$n_pieces, 32100L)
+expect_equal(tok$eos_id, 1L)
+expect_equal(tok$pad_id, 0L)
+expect_equal(tok$unk_id, 2L)
+
+# --- raw encodings (no padding/truncation) --------------------------------------
+
+for (case in cases$cases) {
+ expected <- as.integer(unlist(case$ids))
+ got <- encode_unigram(tok, case$text, max_length = NULL, pad = FALSE)
+ expect_equal(got$input_ids[[1]], expected,
+ info = sprintf("text: %s", substr(case$text, 1, 60)))
+}
+
+# --- padding + truncation at max_length -------------------------------------------
+
+for (case in cases$padded) {
+ got <- encode_unigram(tok, case$text, max_length = case$max_length,
+ pad = TRUE)
+ expect_equal(as.integer(got$input_ids[1, ]), as.integer(unlist(case$ids)),
+ info = sprintf("padded text: %s", substr(case$text, 1, 60)))
+ expect_equal(as.integer(got$attention_mask[1, ]),
+ as.integer(unlist(case$mask)))
+}
+
+# --- batch shape -------------------------------------------------------------------
+
+batch <- encode_unigram(tok, c("a photo of a cat", "hello"),
+ max_length = 32L)
+expect_equal(dim(batch$input_ids), c(2L, 32L))
+expect_equal(dim(batch$attention_mask), c(2L, 32L))
diff --git a/inst/tinytest/test_txt2img_flux.R b/inst/tinytest/test_txt2img_flux.R
new file mode 100644
index 0000000..8e3365f
--- /dev/null
+++ b/inst/tinytest/test_txt2img_flux.R
@@ -0,0 +1,80 @@
+# End-to-end smoke test for the FLUX pipeline wiring on the CPU with
+# tiny random-init components: pack -> denoise (2 steps) -> unpack ->
+# decode. Verifies shapes, finiteness, and the phase plumbing - numeric
+# quality comes from the per-component parity tests.
+
+if (!requireNamespace("torch", quietly = TRUE) || !torch::torch_is_installed()) {
+ exit_file("torch not fully installed")
+}
+
+library(diffuseR)
+
+torch::torch_manual_seed(7)
+
+transformer <- flux_transformer(
+ in_channels = 64L,
+ num_layers = 1L,
+ num_single_layers = 1L,
+ attention_head_dim = 8L,
+ num_attention_heads = 2L,
+ joint_attention_dim = 16L,
+ pooled_projection_dim = 12L,
+ axes_dims_rope = c(2L, 2L, 4L)
+)
+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,
+ format = "full",
+ attn_chunk = NULL,
+ config = list(in_channels = 64L),
+ scheduler_shift = 1.0,
+ vae_scaling_factor = 0.3611,
+ vae_shift_factor = 0.1159
+ ),
+ class = "flux_pipeline"
+)
+
+res <- txt2img_flux(
+ "tiny smoke test",
+ pipeline = pipeline,
+ width = 64L, height = 64L,
+ num_inference_steps = 2L,
+ seed = 42L,
+ prompt_embeds = torch::torch_randn(1L, 7L, 16L),
+ pooled_prompt_embeds = torch::torch_randn(1L, 12L),
+ 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, "flux1-schnell")
+
+# Same seed reproduces the same image
+res2 <- txt2img_flux(
+ "tiny smoke test",
+ pipeline = pipeline,
+ width = 64L, height = 64L,
+ num_inference_steps = 2L,
+ seed = 42L,
+ prompt_embeds = torch::torch_randn(1L, 7L, 16L)$zero_()$add(0.1),
+ pooled_prompt_embeds = torch::torch_randn(1L, 12L)$zero_()$add(0.1),
+ save_file = FALSE,
+ verbose = FALSE
+)
+expect_true(is.finite(max(abs(res2$image))))
diff --git a/man/VAEAttentionBlock.Rd b/man/VAEAttentionBlock.Rd
index 7f1397f..57c76f2 100644
--- a/man/VAEAttentionBlock.Rd
+++ b/man/VAEAttentionBlock.Rd
@@ -3,7 +3,7 @@
\alias{VAEAttentionBlock}
\title{VAE Attention Block}
\usage{
-VAEAttentionBlock(channels)
+VAEAttentionBlock(channels, norm_groups = 32)
}
\arguments{
\item{channels}{Number of channels}
diff --git a/man/VAEMidBlock.Rd b/man/VAEMidBlock.Rd
index f7356df..76bf961 100644
--- a/man/VAEMidBlock.Rd
+++ b/man/VAEMidBlock.Rd
@@ -3,7 +3,7 @@
\alias{VAEMidBlock}
\title{VAE Mid Block}
\usage{
-VAEMidBlock(channels)
+VAEMidBlock(channels, norm_groups = 32)
}
\arguments{
\item{channels}{Number of channels}
diff --git a/man/VAEResnetBlock.Rd b/man/VAEResnetBlock.Rd
index 81ac5dd..92b5a2b 100644
--- a/man/VAEResnetBlock.Rd
+++ b/man/VAEResnetBlock.Rd
@@ -3,7 +3,7 @@
\alias{VAEResnetBlock}
\title{VAE ResNet Block}
\usage{
-VAEResnetBlock(in_channels, out_channels)
+VAEResnetBlock(in_channels, out_channels, norm_groups = 32)
}
\arguments{
\item{in_channels}{Input channels}
diff --git a/man/VAEUpBlock.Rd b/man/VAEUpBlock.Rd
index 5af5ba3..5660e67 100644
--- a/man/VAEUpBlock.Rd
+++ b/man/VAEUpBlock.Rd
@@ -3,7 +3,8 @@
\alias{VAEUpBlock}
\title{VAE Up Block}
\usage{
-VAEUpBlock(in_channels, out_channels, num_resnets = 3, add_upsample = TRUE)
+VAEUpBlock(in_channels, out_channels, num_resnets = 3, add_upsample = TRUE,
+ norm_groups = 32)
}
\arguments{
\item{in_channels}{Input channels}
diff --git a/man/checkpoint_flux.Rd b/man/checkpoint_flux.Rd
new file mode 100644
index 0000000..1ededcc
--- /dev/null
+++ b/man/checkpoint_flux.Rd
@@ -0,0 +1,14 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{checkpoint_flux}
+\alias{checkpoint_flux}
+\title{FLUX Checkpoint Readers}
+\description{
+FLUX transformers ship in the diffusers layout: a directory with
+\code{config.json}, one or more \code{diffusion_pytorch_model*.safetensors}
+shards, and (when sharded) a
+\code{diffusion_pytorch_model.safetensors.index.json} weight map.
+These helpers open that layout behind the same checkpoint interface as
+\code{\link{ltx23_open_checkpoint}}, so the LTX group loaders and
+quantization machinery work unchanged. FLUX module names mirror the
+checkpoint keys 1:1 - no key mapping is needed.
+}
diff --git a/man/clip_pooled_output.Rd b/man/clip_pooled_output.Rd
new file mode 100644
index 0000000..23b87c4
--- /dev/null
+++ b/man/clip_pooled_output.Rd
@@ -0,0 +1,23 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{clip_pooled_output}
+\alias{clip_pooled_output}
+\title{Pooled CLIP output at the EOS position}
+\usage{
+clip_pooled_output(hidden_states, input_ids)
+}
+\arguments{
+\item{hidden_states}{Final-LN hidden states [B, S, D] from
+\code{\link{text_encoder_native}} (with \code{apply_final_ln = TRUE}).}
+
+\item{input_ids}{Token ids [B, S] (0-based, as fed to the encoder).}
+}
+\value{
+Tensor [B, D].
+}
+\description{
+The HF CLIPTextModel pooler_output: the final-layer-norm hidden state
+at the EOS token position, located by argmax over the token ids (EOS
+is the highest id in the CLIP vocab, and causal attention makes any
+padding after it irrelevant). No text projection is applied - this is
+what FLUX uses as pooled_projections.
+}
diff --git a/man/dit_flux.Rd b/man/dit_flux.Rd
new file mode 100644
index 0000000..7c8a8c8
--- /dev/null
+++ b/man/dit_flux.Rd
@@ -0,0 +1,13 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{dit_flux}
+\alias{dit_flux}
+\title{FLUX Transformer (MMDiT)}
+\description{
+Fresh R port of FluxTransformer2DModel from the diffusers reference
+implementation (Apache-2.0,
+src/diffusers/models/transformers/transformer_flux.py). The module
+tree mirrors the diffusers state-dict keys 1:1, so checkpoints load
+without remapping. FLUX.1-schnell has no guidance embedder
+(guidance_embeds = FALSE); the guidance-distilled dev variant is not
+implemented.
+}
diff --git a/man/dit_flux_modules.Rd b/man/dit_flux_modules.Rd
new file mode 100644
index 0000000..bc73d8a
--- /dev/null
+++ b/man/dit_flux_modules.Rd
@@ -0,0 +1,13 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{dit_flux_modules}
+\alias{dit_flux_modules}
+\title{FLUX Transformer Building Blocks}
+\description{
+Fresh R port of the FLUX MMDiT blocks from the diffusers reference
+implementation (Apache-2.0,
+src/diffusers/models/transformers/transformer_flux.py and
+src/diffusers/models/normalization.py). Module field names mirror the
+diffusers state-dict keys 1:1 so checkpoints load without remapping.
+Reuses the LTX primitives \code{ltx23_rms_norm}, \code{.ltx23_sdpa}
+and \code{ltx23_feed_forward}.
+}
diff --git a/man/dot-ltx23_jit_run_stack.Rd b/man/dot-ltx23_jit_run_stack.Rd
index 4d7f3e1..4c53f2d 100644
--- a/man/dot-ltx23_jit_run_stack.Rd
+++ b/man/dot-ltx23_jit_run_stack.Rd
@@ -10,7 +10,8 @@
temb_ca_audio_gate, temb_prompt, temb_prompt_audio,
video_rotary_emb, audio_rotary_emb, ca_video_rotary_emb,
ca_audio_rotary_emb, encoder_attention_mask = NULL,
- audio_encoder_attention_mask = NULL)
+ audio_encoder_attention_mask = NULL,
+ cond_token_index = NULL)
}
\value{
list(hidden_states, audio_hidden_states)
diff --git a/man/download_flux.Rd b/man/download_flux.Rd
new file mode 100644
index 0000000..d4a0fc4
--- /dev/null
+++ b/man/download_flux.Rd
@@ -0,0 +1,9 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{download_flux}
+\alias{download_flux}
+\title{Download and Prepare FLUX.1-schnell Weights}
+\description{
+Downloads FLUX.1-schnell from HuggingFace (weights Apache-2.0, but
+the repo is gated behind a license click-through) and quantizes the
+12B transformer to a local NF4 (~7 GB) or fp8 (~12 GB) artifact.
+}
diff --git a/man/download_flux1.Rd b/man/download_flux1.Rd
new file mode 100644
index 0000000..4e68488
--- /dev/null
+++ b/man/download_flux1.Rd
@@ -0,0 +1,33 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{download_flux1}
+\alias{download_flux1}
+\title{Download FLUX.1-schnell and build the quantized artifact}
+\usage{
+download_flux1(quantize = TRUE, precision = c("nf4", "fp8"), output_dir = NULL,
+ text_encoders = TRUE, verbose = TRUE)
+}
+\arguments{
+\item{quantize}{Logical. Build the quantized artifact after
+downloading.}
+
+\item{precision}{"nf4" (~7 GB, GPU-resident on 16 GB cards) or
+"fp8" (~12 GB, CPU-resident, streamed; near-bf16 quality).}
+
+\item{output_dir}{Directory for the quantized artifact.}
+
+\item{text_encoders}{Logical. Also fetch the CLIP + T5 text encoders,
+tokenizer, VAE, and scheduler config (~10 GB).}
+
+\item{verbose}{Logical.}
+}
+\value{
+Invisibly, a list with \code{transformer_dir},
+ \code{artifact_dir}, and \code{support} (named file paths).
+}
+\description{
+Skips work already done: a valid quantized manifest short-circuits
+the transformer download; cached files are not re-fetched. Needs
+\code{HF_TOKEN} set for the gated repo (see the error message it
+raises without one). The bf16 transformer source (~24 GB in the
+HuggingFace cache) may be deleted after quantization.
+}
diff --git a/man/encode_unigram.Rd b/man/encode_unigram.Rd
new file mode 100644
index 0000000..0bd0679
--- /dev/null
+++ b/man/encode_unigram.Rd
@@ -0,0 +1,32 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{encode_unigram}
+\alias{encode_unigram}
+\title{Encode text with a Unigram tokenizer}
+\usage{
+encode_unigram(tokenizer, texts, max_length = 256L, add_eos = TRUE, pad = TRUE)
+}
+\arguments{
+\item{tokenizer}{A \code{\link{unigram_tokenizer}}.}
+
+\item{texts}{Character vector of prompts.}
+
+\item{max_length}{Integer. Fixed sequence length (NULL for no
+truncation/padding).}
+
+\item{add_eos}{Logical. Append the EOS token.}
+
+\item{pad}{Logical. Right-pad to \code{max_length}.}
+}
+\value{
+List with \code{input_ids} and \code{attention_mask}, each an
+ integer matrix [length(texts), max_length] (or ragged lists when
+ \code{max_length} is NULL). Ids are 0-based (HuggingFace
+ convention); add 1 for R torch embedding lookups.
+}
+\description{
+Normalizes (strip-right, multi-space collapse, control whitespace to
+space), applies the Metaspace pre-tokenizer, segments each pre-token
+by Viterbi over the Unigram scores, fuses consecutive unknowns, and
+appends EOS. T5 semantics: right padding with \code{} (id 0),
+truncation to \code{max_length - 1} before the EOS.
+}
diff --git a/man/encode_with_t5.Rd b/man/encode_with_t5.Rd
new file mode 100644
index 0000000..d01c4c4
--- /dev/null
+++ b/man/encode_with_t5.Rd
@@ -0,0 +1,27 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{encode_with_t5}
+\alias{encode_with_t5}
+\title{Encode prompts with the T5 encoder}
+\usage{
+encode_with_t5(prompts, model, tokenizer, max_sequence_length = 256L,
+ device = NULL)
+}
+\arguments{
+\item{prompts}{Character vector.}
+
+\item{model}{A \code{\link{t5_encoder}}.}
+
+\item{tokenizer}{A \code{\link{unigram_tokenizer}}.}
+
+\item{max_sequence_length}{Integer. Fixed token length (schnell: 256).}
+
+\item{device}{Device for the input ids (defaults to the model's).}
+}
+\value{
+Tensor [length(prompts), max_sequence_length, d_model].
+}
+\description{
+Tokenizes with \code{\link{encode_unigram}} (right padding to
+\code{max_sequence_length}) and runs the encoder. Matching the FLUX
+reference pipeline, no attention mask is used.
+}
diff --git a/man/flux_ada_layer_norm_continuous.Rd b/man/flux_ada_layer_norm_continuous.Rd
new file mode 100644
index 0000000..1504a4c
--- /dev/null
+++ b/man/flux_ada_layer_norm_continuous.Rd
@@ -0,0 +1,18 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_ada_layer_norm_continuous}
+\alias{flux_ada_layer_norm_continuous}
+\title{FLUX continuous adaLN (final norm)}
+\usage{
+flux_ada_layer_norm_continuous(dim, cond_dim = dim)
+}
+\arguments{
+\item{dim}{Integer. Model dimension.}
+
+\item{cond_dim}{Integer. Conditioning embedding dimension.}
+}
+\description{
+Scale/shift conditioning of the final norm. Note the chunk order:
+scale first, then shift (the reverse of adaLN-Zero). Reference:
+diffusers AdaLayerNormContinuous as used by FLUX norm_out
+(elementwise_affine = FALSE, eps = 1e-6).
+}
diff --git a/man/flux_ada_layer_norm_zero.Rd b/man/flux_ada_layer_norm_zero.Rd
new file mode 100644
index 0000000..32806e1
--- /dev/null
+++ b/man/flux_ada_layer_norm_zero.Rd
@@ -0,0 +1,19 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_ada_layer_norm_zero}
+\alias{flux_ada_layer_norm_zero}
+\title{FLUX adaLN-Zero modulation (double-stream)}
+\usage{
+flux_ada_layer_norm_zero(dim)
+}
+\arguments{
+\item{dim}{Integer. Model dimension.}
+}
+\value{
+Module whose forward(x, emb) returns
+ \code{list(x_norm, gate_msa, shift_mlp, scale_mlp, gate_mlp)}.
+}
+\description{
+Projects the conditioning embedding to six modulation vectors and
+returns the msa-modulated input plus the remaining parameters.
+Reference: diffusers AdaLayerNormZero.
+}
diff --git a/man/flux_ada_layer_norm_zero_single.Rd b/man/flux_ada_layer_norm_zero_single.Rd
new file mode 100644
index 0000000..1aea97a
--- /dev/null
+++ b/man/flux_ada_layer_norm_zero_single.Rd
@@ -0,0 +1,17 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_ada_layer_norm_zero_single}
+\alias{flux_ada_layer_norm_zero_single}
+\title{FLUX adaLN-Zero modulation (single-stream)}
+\usage{
+flux_ada_layer_norm_zero_single(dim)
+}
+\arguments{
+\item{dim}{Integer. Model dimension.}
+}
+\value{
+Module whose forward(x, emb) returns \code{list(x_norm, gate)}.
+}
+\description{
+Three modulation vectors: shift, scale, gate. Reference: diffusers
+AdaLayerNormZeroSingle.
+}
diff --git a/man/flux_apply_rotary_emb.Rd b/man/flux_apply_rotary_emb.Rd
new file mode 100644
index 0000000..1b92d27
--- /dev/null
+++ b/man/flux_apply_rotary_emb.Rd
@@ -0,0 +1,23 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_apply_rotary_emb}
+\alias{flux_apply_rotary_emb}
+\title{Apply FLUX rotary embeddings to a per-head tensor}
+\usage{
+flux_apply_rotary_emb(x, freqs)
+}
+\arguments{
+\item{x}{Tensor of shape [B, H, S, D] (per-head layout).}
+
+\item{freqs}{List of two tensors (cos, sin), each [S, D], from
+\code{flux_pos_embed}.}
+}
+\value{
+Tensor with the same shape and dtype as \code{x}.
+}
+\description{
+Rotates adjacent element pairs of the last dimension:
+\code{out = x * cos + rotate_half(x) * sin} with pairs interleaved
+(elements 1,2 form the first complex pair). Math in float32, result
+cast back to the input dtype. Reference: apply_rotary_emb with
+use_real_unbind_dim = -1.
+}
diff --git a/man/flux_attention.Rd b/man/flux_attention.Rd
new file mode 100644
index 0000000..9ba74f4
--- /dev/null
+++ b/man/flux_attention.Rd
@@ -0,0 +1,30 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_attention}
+\alias{flux_attention}
+\title{FLUX joint attention}
+\usage{
+flux_attention(query_dim, heads, dim_head, added_kv = FALSE, pre_only = FALSE,
+ eps = 1e-06)
+}
+\arguments{
+\item{query_dim}{Integer. Model dimension.}
+
+\item{heads}{Integer. Number of attention heads.}
+
+\item{dim_head}{Integer. Per-head dimension.}
+
+\item{added_kv}{Logical. Add text-stream projections (double blocks).}
+
+\item{pre_only}{Logical. Skip the output projection (single blocks).}
+
+\item{eps}{Numeric. RMS norm epsilon.}
+}
+\description{
+Multi-head attention with per-head RMS q/k norms and rotary position
+embeddings. With \code{added_kv = TRUE} (double-stream blocks) the
+text stream gets its own q/k/v projections and both streams attend
+jointly (text tokens first); the outputs are split back and projected
+per stream. With \code{pre_only = TRUE} (single-stream blocks) there
+is no output projection. Reference: diffusers FluxAttention +
+FluxAttnProcessor.
+}
diff --git a/man/flux_double_block.Rd b/man/flux_double_block.Rd
new file mode 100644
index 0000000..534e11a
--- /dev/null
+++ b/man/flux_double_block.Rd
@@ -0,0 +1,24 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_double_block}
+\alias{flux_double_block}
+\title{FLUX double-stream (MMDiT) transformer block}
+\usage{
+flux_double_block(dim, num_attention_heads, attention_head_dim)
+}
+\arguments{
+\item{dim}{Integer. Model dimension.}
+
+\item{num_attention_heads}{Integer. Attention heads.}
+
+\item{attention_head_dim}{Integer. Per-head dimension.}
+}
+\value{
+Module whose forward(hidden_states, encoder_hidden_states,
+ temb, image_rotary_emb) returns
+ \code{list(encoder_hidden_states, hidden_states)}.
+}
+\description{
+Image and text streams each get adaLN-Zero modulation and a
+feed-forward; attention is joint across both streams. Reference:
+diffusers FluxTransformerBlock.
+}
diff --git a/man/flux_is_quant_key.Rd b/man/flux_is_quant_key.Rd
new file mode 100644
index 0000000..a06d2c9
--- /dev/null
+++ b/man/flux_is_quant_key.Rd
@@ -0,0 +1,16 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_is_quant_key}
+\alias{flux_is_quant_key}
+\title{Test whether a FLUX key is in the quantization cast set}
+\usage{
+flux_is_quant_key(key)
+}
+\arguments{
+\item{key}{Character vector of parameter names (diffusers-style).}
+}
+\value{
+Logical vector.
+}
+\description{
+Test whether a FLUX key is in the quantization cast set
+}
diff --git a/man/flux_load_pipeline.Rd b/man/flux_load_pipeline.Rd
new file mode 100644
index 0000000..cef84a9
--- /dev/null
+++ b/man/flux_load_pipeline.Rd
@@ -0,0 +1,38 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_load_pipeline}
+\alias{flux_load_pipeline}
+\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)
+}
+\arguments{
+\item{model_dir}{Quantized artifact directory (default: the
+\code{download_flux1} location for \code{precision}), or a raw
+diffusers transformer directory for full-precision loading.}
+
+\item{device}{Character. Compute device.}
+
+\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{attn_chunk}{Integer or NULL. Attention query-chunk override.}
+
+\item{phase_offload}{Logical. One GPU tenant per phase.}
+
+\item{verbose}{Logical.}
+}
+\value{
+A \code{flux_pipeline} list.
+}
+\description{
+Loads the quantized transformer artifact plus the VAE decoder, CLIP
+and T5 text encoders, tokenizer, and scheduler config (from the
+HuggingFace cache populated by \code{\link{download_flux1}}).
+Components load to the CPU when \code{phase_offload} is on and move
+to the GPU only for their phase of the generation.
+}
diff --git a/man/flux_load_transformer.Rd b/man/flux_load_transformer.Rd
new file mode 100644
index 0000000..112cb8f
--- /dev/null
+++ b/man/flux_load_transformer.Rd
@@ -0,0 +1,46 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_load_transformer}
+\alias{flux_load_transformer}
+\title{Load a FLUX transformer from any checkpoint format}
+\usage{
+flux_load_transformer(ckpt, device = "cuda", dtype = "bfloat16", pin = TRUE,
+ verbose = TRUE, ...)
+}
+\arguments{
+\item{ckpt}{A checkpoint from \code{\link{flux_open_checkpoint}} or
+\code{\link{flux_open_quantized}}.}
+
+\item{device}{Character. Compute device.}
+
+\item{dtype}{Character. Model dtype ("bfloat16" or "float32"). For
+quantized formats this sets the resident (non-quantized) tensors
+and must match the compute dtype: bfloat16 for GPU compute,
+float32 for CPU compute.}
+
+\item{pin}{Logical. Pin fp8 host memory for faster transfers.}
+
+\item{verbose}{Logical.}
+
+\item{...}{Overrides for \code{\link{flux_transformer}} arguments
+(tiny test configs).}
+}
+\value{
+The loaded \code{flux_transformer} in eval mode.
+}
+\description{
+Builds \code{\link{flux_transformer}} from the checkpoint's embedded
+config and loads the weights. Dispatches on the checkpoint format:
+}
+\details{
+\itemize{
+\item full precision (\code{\link{flux_open_checkpoint}}): weights
+stream into the model in \code{dtype} on \code{device}.
+\item \code{"nf4"} (\code{\link{flux_open_quantized}}): cast-set
+linears become \code{ltx23_nf4_linear}; the whole model (packed
+weights included) moves to \code{device} and stays resident.
+\item \code{"fp8"}: cast-set linears become
+\code{ltx23_fp8_linear}; fp8 weights stay CPU-resident (optionally
+pinned) and stream to \code{device} inside each forward.
+}
+
+}
diff --git a/man/flux_memory_profile.Rd b/man/flux_memory_profile.Rd
new file mode 100644
index 0000000..4d5c7df
--- /dev/null
+++ b/man/flux_memory_profile.Rd
@@ -0,0 +1,19 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_memory_profile}
+\alias{flux_memory_profile}
+\title{Resolve a FLUX memory profile}
+\usage{
+flux_memory_profile(vram_gb = NULL)
+}
+\arguments{
+\item{vram_gb}{Numeric or NULL. Available VRAM; auto-detected when
+NULL (via gpu.ctl or nvidia-smi).}
+}
+\value{
+List with \code{name}, \code{precision} ("nf4"/"fp8"),
+ \code{attn_chunk}, \code{text_device}, \code{phase_offload}, and
+ \code{max_pixels} (largest validated image area).
+}
+\description{
+Resolve a FLUX memory profile
+}
diff --git a/man/flux_open_checkpoint.Rd b/man/flux_open_checkpoint.Rd
new file mode 100644
index 0000000..437cd09
--- /dev/null
+++ b/man/flux_open_checkpoint.Rd
@@ -0,0 +1,22 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_open_checkpoint}
+\alias{flux_open_checkpoint}
+\title{Open a FLUX transformer checkpoint directory}
+\usage{
+flux_open_checkpoint(transformer_dir)
+}
+\arguments{
+\item{transformer_dir}{Directory containing \code{config.json} and the
+\code{diffusion_pytorch_model*.safetensors} file(s).}
+}
+\value{
+An object of class \code{ltx23_checkpoint} (shared checkpoint
+ interface): list with \code{handle$get_tensor}, \code{keys},
+ \code{config}, and \code{path}.
+}
+\description{
+Opens a diffusers-layout transformer directory lazily (headers only).
+Sharded checkpoints are resolved through the index.json weight map;
+single-file checkpoints are opened directly. The transformer
+\code{config.json} is attached as \code{$config}.
+}
diff --git a/man/flux_open_quantized.Rd b/man/flux_open_quantized.Rd
new file mode 100644
index 0000000..7aa34bc
--- /dev/null
+++ b/man/flux_open_quantized.Rd
@@ -0,0 +1,19 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_open_quantized}
+\alias{flux_open_quantized}
+\title{Open a quantized FLUX artifact directory}
+\usage{
+flux_open_quantized(dir)
+}
+\arguments{
+\item{dir}{The quantized artifact directory (with manifest.json).}
+}
+\value{
+An \code{ltx23_checkpoint} with \code{$format} set.
+}
+\description{
+Opens the sharded NF4/fp8 artifact written by
+\code{\link{flux_quantize}} through the shared checkpoint interface.
+The manifest's embedded transformer config and \code{format} ride
+along, so \code{\link{flux_load_transformer}} needs nothing else.
+}
diff --git a/man/flux_pack_latents.Rd b/man/flux_pack_latents.Rd
new file mode 100644
index 0000000..7af3997
--- /dev/null
+++ b/man/flux_pack_latents.Rd
@@ -0,0 +1,17 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_pack_latents}
+\alias{flux_pack_latents}
+\title{Pack FLUX latents into a patch sequence}
+\usage{
+flux_pack_latents(latents)
+}
+\arguments{
+\item{latents}{Tensor of shape [B, C, H, W]; H and W must be even.}
+}
+\value{
+Tensor of shape [B, (H/2) * (W/2), C * 4].
+}
+\description{
+Packs a [B, C, H, W] latent into 2x2 patches, giving a sequence
+[B, (H/2) * (W/2), C * 4]. Reference: FluxPipeline._pack_latents.
+}
diff --git a/man/flux_pos_embed.Rd b/man/flux_pos_embed.Rd
new file mode 100644
index 0000000..6d3ae1e
--- /dev/null
+++ b/man/flux_pos_embed.Rd
@@ -0,0 +1,26 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_pos_embed}
+\alias{flux_pos_embed}
+\title{Compute FLUX rotary frequencies from position ids}
+\usage{
+flux_pos_embed(ids, axes_dim = c(16L, 56L, 56L), theta = 10000)
+}
+\arguments{
+\item{ids}{Tensor of shape [S, 3]: concatenated text ids (all zero)
+and image ids from \code{flux_prepare_latent_image_ids}.}
+
+\item{axes_dim}{Integer vector of per-axis rotary dims; must sum to
+the attention head dim. FLUX uses c(16, 56, 56).}
+
+\item{theta}{Numeric. RoPE base frequency.}
+}
+\value{
+List of two tensors (cos, sin), each [S, sum(axes_dim)],
+ float32, on the device of \code{ids}.
+}
+\description{
+Per-axis 1D rotary frequencies (interleaved-real convention), computed
+in float64 on CPU and concatenated over the axes. Reference:
+FluxPosEmbed with get_1d_rotary_pos_embed(repeat_interleave_real=TRUE,
+use_real=TRUE, freqs_dtype=float64).
+}
diff --git a/man/flux_prepare_latent_image_ids.Rd b/man/flux_prepare_latent_image_ids.Rd
new file mode 100644
index 0000000..5d832b6
--- /dev/null
+++ b/man/flux_prepare_latent_image_ids.Rd
@@ -0,0 +1,22 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_prepare_latent_image_ids}
+\alias{flux_prepare_latent_image_ids}
+\title{Build FLUX latent image position ids}
+\usage{
+flux_prepare_latent_image_ids(height, width, device = "cpu")
+}
+\arguments{
+\item{height}{Integer. Packed grid height (latent height / 2).}
+
+\item{width}{Integer. Packed grid width (latent width / 2).}
+
+\item{device}{Device for the resulting tensor.}
+}
+\value{
+Float tensor of shape [height * width, 3].
+}
+\description{
+Position ids over the packed latent grid (latent height/2 x width/2).
+Channel 1 is always zero, channel 2 holds the row index, channel 3 the
+column index. Reference: FluxPipeline._prepare_latent_image_ids.
+}
diff --git a/man/flux_quantize.Rd b/man/flux_quantize.Rd
new file mode 100644
index 0000000..4f97593
--- /dev/null
+++ b/man/flux_quantize.Rd
@@ -0,0 +1,34 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_quantize}
+\alias{flux_quantize}
+\title{Quantize a FLUX transformer to NF4 or fp8 shards}
+\usage{
+flux_quantize(transformer_dir, output_dir = NULL, format = c("nf4", "fp8"),
+ shard_bytes = 4e+09, force = FALSE, verbose = TRUE)
+}
+\arguments{
+\item{transformer_dir}{Source diffusers transformer directory.}
+
+\item{output_dir}{Output directory for shards + manifest (default:
+the per-format location under \code{tools::R_user_dir}).}
+
+\item{format}{"nf4" or "fp8".}
+
+\item{shard_bytes}{Numeric. Approximate shard size.}
+
+\item{force}{Logical. Re-quantize even if a valid manifest exists.}
+
+\item{verbose}{Logical.}
+}
+\value{
+Invisibly, the manifest list.
+}
+\description{
+Streams the bf16 diffusers checkpoint tensor by tensor. Cast-set
+weights (see \code{\link{flux_is_quant_key}}) are stored as NF4
+(packed uint8 + \code{_absmax} float32 blocks) or as
+float8_e4m3fn with an absmax/448 per-tensor \code{_scale};
+everything else is copied through unchanged. The manifest embeds the
+transformer config, so the source checkpoint is not needed again
+after quantization.
+}
diff --git a/man/flux_single_block.Rd b/man/flux_single_block.Rd
new file mode 100644
index 0000000..22aee02
--- /dev/null
+++ b/man/flux_single_block.Rd
@@ -0,0 +1,28 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_single_block}
+\alias{flux_single_block}
+\title{FLUX single-stream transformer block}
+\usage{
+flux_single_block(dim, num_attention_heads, attention_head_dim, mlp_ratio = 4)
+}
+\arguments{
+\item{dim}{Integer. Model dimension.}
+
+\item{num_attention_heads}{Integer. Attention heads.}
+
+\item{attention_head_dim}{Integer. Per-head dimension.}
+
+\item{mlp_ratio}{Numeric. MLP hidden dim multiplier.}
+}
+\value{
+Module whose forward(hidden_states, temb, image_rotary_emb)
+ returns the joint hidden states.
+}
+\description{
+Parallel attention + MLP over the joint [text; image] sequence with a
+shared gate: \code{x + gate * proj_out(cat(attn, gelu(mlp)))}. The
+reference concatenates the streams inside every block and splits after;
+here the caller concatenates once before the single-block stack, which
+is numerically identical. Reference: diffusers
+FluxSingleTransformerBlock.
+}
diff --git a/man/flux_transformer.Rd b/man/flux_transformer.Rd
new file mode 100644
index 0000000..86b5cd9
--- /dev/null
+++ b/man/flux_transformer.Rd
@@ -0,0 +1,44 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_transformer}
+\alias{flux_transformer}
+\title{FLUX transformer model}
+\usage{
+flux_transformer(in_channels = 64L, num_layers = 19L, num_single_layers = 38L,
+ attention_head_dim = 128L, num_attention_heads = 24L,
+ joint_attention_dim = 4096L, pooled_projection_dim = 768L,
+ axes_dims_rope = c(16L, 56L, 56L), out_channels = NULL)
+}
+\arguments{
+\item{in_channels}{Integer. Packed latent channels (64).}
+
+\item{num_layers}{Integer. Double-stream block count.}
+
+\item{num_single_layers}{Integer. Single-stream block count.}
+
+\item{attention_head_dim}{Integer. Per-head dimension.}
+
+\item{num_attention_heads}{Integer. Attention heads.}
+
+\item{joint_attention_dim}{Integer. T5 embedding dim (4096).}
+
+\item{pooled_projection_dim}{Integer. CLIP pooled dim (768).}
+
+\item{axes_dims_rope}{Integer vector. Per-axis rotary dims.}
+
+\item{out_channels}{Integer or NULL. Output channels (defaults to
+\code{in_channels}).}
+}
+\value{
+Module whose forward(hidden_states, encoder_hidden_states,
+ pooled_projections, timestep, image_rotary_emb) returns the
+ predicted velocity for the image tokens [B, S_img, out_channels].
+ \code{timestep} is in sigma space (0-1); it is scaled by 1000
+ internally, matching the reference.
+}
+\description{
+19 double-stream (MMDiT) blocks followed by 38 single-stream blocks
+over the joint [text; image] sequence, with adaLN-Zero conditioning on
+timestep + pooled CLIP text. Rotary embeddings are precomputed by the
+caller with \code{flux_pos_embed} (they are static across denoise
+steps). Defaults are the FLUX.1-schnell configuration.
+}
diff --git a/man/flux_unpack_latents.Rd b/man/flux_unpack_latents.Rd
new file mode 100644
index 0000000..93a10d5
--- /dev/null
+++ b/man/flux_unpack_latents.Rd
@@ -0,0 +1,23 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{flux_unpack_latents}
+\alias{flux_unpack_latents}
+\title{Unpack a FLUX patch sequence back into latents}
+\usage{
+flux_unpack_latents(latents, height, width, vae_scale_factor = 8L)
+}
+\arguments{
+\item{latents}{Tensor of shape [B, S, C_packed].}
+
+\item{vae_scale_factor}{Integer. Spatial downsampling of the VAE (8).}
+
+\item{height,width}{Integers. Image height/width in pixels.}
+}
+\value{
+Tensor of shape [B, C_packed / 4, height / 8, width / 8].
+}
+\description{
+Inverse of \code{flux_pack_latents}. Height and width are the target
+image dimensions in pixels; the latent grid is derived via the VAE
+scale factor and the 2x2 patch size. Reference:
+FluxPipeline._unpack_latents.
+}
diff --git a/man/load_decoder_safetensors.Rd b/man/load_decoder_safetensors.Rd
new file mode 100644
index 0000000..bff9fe3
--- /dev/null
+++ b/man/load_decoder_safetensors.Rd
@@ -0,0 +1,25 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{load_decoder_safetensors}
+\alias{load_decoder_safetensors}
+\title{Load HF safetensors VAE weights into the native decoder}
+\usage{
+load_decoder_safetensors(native_decoder, path, verbose = TRUE)
+}
+\arguments{
+\item{native_decoder}{Native VAE decoder module}
+
+\item{path}{Path to the VAE .safetensors file (or a directory
+containing diffusion_pytorch_model.safetensors)}
+
+\item{verbose}{Print loading progress}
+}
+\value{
+The native decoder with loaded weights (invisibly)
+}
+\description{
+Loads the decoder half of a diffusers AutoencoderKL safetensors file
+(e.g. FLUX.1-schnell's \code{vae/diffusion_pytorch_model.safetensors}).
+Keys under \code{decoder.} map to the native module 1:1; encoder and
+quant-conv keys are skipped (the FLUX VAE has no quant convs, and
+txt2img needs no encoder).
+}
diff --git a/man/load_t5_text_encoder.Rd b/man/load_t5_text_encoder.Rd
new file mode 100644
index 0000000..ff34b0c
--- /dev/null
+++ b/man/load_t5_text_encoder.Rd
@@ -0,0 +1,29 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{load_t5_text_encoder}
+\alias{load_t5_text_encoder}
+\title{Load a T5 encoder from a transformers directory}
+\usage{
+load_t5_text_encoder(model_path, device = "cpu", dtype = "float32",
+ verbose = TRUE, ...)
+}
+\arguments{
+\item{model_path}{Directory with \code{config.json} and
+\code{model*.safetensors} (FLUX.1-schnell's \code{text_encoder_2}).}
+
+\item{device}{Character. Target device.}
+
+\item{dtype}{Character. "float32" (CPU default; T5 overflows in
+float16) or "bfloat16".}
+
+\item{verbose}{Logical.}
+
+\item{...}{Overrides for \code{\link{t5_encoder}} arguments.}
+}
+\value{
+The loaded \code{t5_encoder} in eval mode.
+}
+\description{
+Streams the (possibly sharded) safetensors weights into
+\code{\link{t5_encoder}}, stripping the \code{encoder.} key prefix
+and aliasing \code{embed_tokens} to the shared embedding.
+}
diff --git a/man/load_text_encoder_safetensors.Rd b/man/load_text_encoder_safetensors.Rd
new file mode 100644
index 0000000..0cb6b4b
--- /dev/null
+++ b/man/load_text_encoder_safetensors.Rd
@@ -0,0 +1,23 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{load_text_encoder_safetensors}
+\alias{load_text_encoder_safetensors}
+\title{Load HF safetensors weights into the native CLIP text encoder}
+\usage{
+load_text_encoder_safetensors(native_encoder, path, verbose = TRUE)
+}
+\arguments{
+\item{native_encoder}{Native text encoder module}
+
+\item{path}{Path to model.safetensors (or a directory containing it)}
+
+\item{verbose}{Print loading progress}
+}
+\value{
+The native encoder with loaded weights (invisibly)
+}
+\description{
+Loads a HuggingFace CLIPTextModel \code{model.safetensors} (e.g.
+FLUX.1-schnell's \code{text_encoder}) into
+\code{\link{text_encoder_native}}, reusing the TorchScript key remaps
+minus the export prefixes.
+}
diff --git a/man/memory_flux.Rd b/man/memory_flux.Rd
new file mode 100644
index 0000000..4c79371
--- /dev/null
+++ b/man/memory_flux.Rd
@@ -0,0 +1,10 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{memory_flux}
+\alias{memory_flux}
+\title{FLUX Memory Profiles}
+\description{
+VRAM-based execution profiles for the FLUX.1-schnell pipeline,
+following the LTX-2.3 profile pattern. The 12B transformer runs NF4
+(~7 GB, GPU-resident) or fp8 (~12 GB, CPU-resident and streamed);
+the T5-XXL text encoder runs float32 on the CPU by default.
+}
diff --git a/man/quantize_flux.Rd b/man/quantize_flux.Rd
new file mode 100644
index 0000000..09351c1
--- /dev/null
+++ b/man/quantize_flux.Rd
@@ -0,0 +1,12 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{quantize_flux}
+\alias{quantize_flux}
+\title{FLUX Transformer Quantization and Loading}
+\description{
+Quantize the 12B FLUX transformer to NF4 (~7 GB, GPU-resident on
+16 GB cards) or fp8 (~12 GB, CPU-resident and streamed per forward),
+and load any format back into \code{\link{flux_transformer}}. Reuses
+the LTX-2.3 quantization machinery (\code{ltx23_nf4_quantize},
+\code{ltx23_nf4_linear}, \code{ltx23_fp8_linear}); only the cast set
+and the diffusers directory layout are FLUX-specific.
+}
diff --git a/man/rope_flux.Rd b/man/rope_flux.Rd
new file mode 100644
index 0000000..a337e96
--- /dev/null
+++ b/man/rope_flux.Rd
@@ -0,0 +1,14 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{rope_flux}
+\alias{rope_flux}
+\title{FLUX Rotary Positional Embeddings}
+\description{
+Fresh R port of the FLUX rotary positional embedding scheme from the
+diffusers reference implementation (Apache-2.0,
+src/diffusers/models/transformers/transformer_flux.py FluxPosEmbed and
+src/diffusers/models/embeddings.py get_1d_rotary_pos_embed /
+apply_rotary_emb). FLUX uses the interleaved adjacent-pair convention
+(use_real_unbind_dim = -1) with per-axis frequencies computed in
+float64 and applied in float32. Text tokens carry all-zero ids, so
+they receive the identity rotation.
+}
diff --git a/man/t5_encoder.Rd b/man/t5_encoder.Rd
new file mode 100644
index 0000000..a41cf1b
--- /dev/null
+++ b/man/t5_encoder.Rd
@@ -0,0 +1,24 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{t5_encoder}
+\alias{t5_encoder}
+\title{T5 encoder stack}
+\usage{
+t5_encoder(vocab_size = 32128L, d_model = 4096L, d_kv = 64L, num_heads = 64L,
+ d_ff = 10240L, num_layers = 24L,
+ relative_attention_num_buckets = 32L,
+ relative_attention_max_distance = 128L, layer_norm_epsilon = 1e-06)
+}
+\arguments{
+\item{layer_norm_epsilon}{Numeric.}
+
+\item{vocab_size,d_model,d_kv,num_heads,d_ff,num_layers}{Integers.}
+
+\item{relative_attention_num_buckets,relative_attention_max_distance}{Integers. Relative position bias shape.}
+}
+\value{
+Module whose forward(input_ids) (1-based ids [B, S]) returns
+ the last hidden state [B, S, d_model].
+}
+\description{
+Defaults are the T5-v1.1-XXL configuration used by FLUX.
+}
diff --git a/man/t5_text_encoder.Rd b/man/t5_text_encoder.Rd
new file mode 100644
index 0000000..57c229c
--- /dev/null
+++ b/man/t5_text_encoder.Rd
@@ -0,0 +1,20 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{t5_text_encoder}
+\alias{t5_text_encoder}
+\title{T5 Text Encoder (T5-v1.1)}
+\description{
+Fresh R port of the T5 encoder stack from HuggingFace transformers
+(Apache-2.0, src/transformers/models/t5/modeling_t5.py), as used by
+FLUX's second text encoder (T5-v1.1-XXL: 24 layers, d_model 4096,
+64 heads x d_kv 64, gated-GELU FFN). Distinctives faithfully carried
+over: RMS layer norms (no mean subtraction), no biases anywhere, no
+1/sqrt(d) attention scaling (folded into the weights), and a shared
+relative position bias computed once from block 1's embedding and
+added to every layer's attention logits. Module field names mirror
+the checkpoint keys (minus the \code{encoder.} prefix).
+}
+\details{
+FLUX passes no attention mask - padding tokens attend and are
+attended to - so none is implemented.
+
+}
diff --git a/man/text_encoder_native.Rd b/man/text_encoder_native.Rd
index 2279664..787cfc0 100644
--- a/man/text_encoder_native.Rd
+++ b/man/text_encoder_native.Rd
@@ -5,7 +5,7 @@
\usage{
text_encoder_native(vocab_size = 49408, context_length = 77, embed_dim = 768,
num_layers = 12, num_heads = 12, mlp_dim = 3072,
- apply_final_ln = TRUE)
+ apply_final_ln = TRUE, gelu_type = "tanh")
}
\arguments{
\item{vocab_size}{Vocabulary size (default 49408)}
@@ -22,6 +22,9 @@ text_encoder_native(vocab_size = 49408, context_length = 77, embed_dim = 768,
\item{apply_final_ln}{Whether to apply final layer norm (default TRUE).
Set to FALSE to match TorchScript exports that don't include final LN.}
+
+\item{gelu_type}{GELU variant: "tanh" (matches the TorchScript exports),
+"quick" (HF CLIP ViT-L, used by FLUX), or "exact"}
}
\value{
An nn_module representing the text encoder
diff --git a/man/tokenizer_unigram.Rd b/man/tokenizer_unigram.Rd
new file mode 100644
index 0000000..9896a65
--- /dev/null
+++ b/man/tokenizer_unigram.Rd
@@ -0,0 +1,18 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{tokenizer_unigram}
+\alias{tokenizer_unigram}
+\title{SentencePiece Unigram Tokenizer}
+\description{
+Pure R implementation of HuggingFace tokenizer.json files with a
+Unigram model (SentencePiece), as used by T5 - FLUX's second text
+encoder. Segmentation is Viterbi best-path over the vocab log
+probabilities (Kudo 2018, arXiv:1804.10959). The normalizer and
+Metaspace pre-tokenizer settings are read from the file.
+}
+\details{
+Limitation: the Precompiled charsmap normalizer (NFKC-style unicode
+mapping) is approximated by control-whitespace substitution only;
+ASCII and common latin text tokenizes identically to the reference,
+exotic unicode may differ.
+
+}
diff --git a/man/txt2img.Rd b/man/txt2img.Rd
index d8d9dd8..7eb0c03 100644
--- a/man/txt2img.Rd
+++ b/man/txt2img.Rd
@@ -3,7 +3,7 @@
\alias{txt2img}
\title{Generate an image from a text prompt using a diffusion pipeline}
\usage{
-txt2img(prompt, model_name = c("sd21", "sdxl"), ...)
+txt2img(prompt, model_name = c("sd21", "sdxl", "flux1"), ...)
}
\arguments{
\item{prompt}{A character string prompt describing the image to generate.}
diff --git a/man/txt2img_flux.Rd b/man/txt2img_flux.Rd
new file mode 100644
index 0000000..e5818d8
--- /dev/null
+++ b/man/txt2img_flux.Rd
@@ -0,0 +1,48 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{txt2img_flux}
+\alias{txt2img_flux}
+\title{Generate an image with FLUX.1-schnell}
+\usage{
+txt2img_flux(prompt, pipeline = NULL, width = 1024L, height = 1024L,
+ num_inference_steps = 4L, max_sequence_length = 256L, seed = NULL,
+ prompt_embeds = NULL, pooled_prompt_embeds = NULL,
+ save_file = TRUE, filename = NULL, verbose = TRUE, ...)
+}
+\arguments{
+\item{prompt}{Character. The prompt.}
+
+\item{pipeline}{A \code{flux_pipeline} from
+\code{\link{flux_load_pipeline}}; NULL loads one (passing
+\code{...} through).}
+
+\item{num_inference_steps}{Integer. Denoising steps (schnell: 4).}
+
+\item{max_sequence_length}{Integer. T5 token length (schnell: 256).}
+
+\item{seed}{Integer or NULL. Initial latents are drawn on the CPU, so
+a seed matches a Python diffusers run with a CPU generator.}
+
+\item{save_file}{Logical. Write a PNG.}
+
+\item{filename}{Output path (default derived from the prompt).}
+
+\item{verbose}{Logical.}
+
+\item{...}{Passed to \code{\link{flux_load_pipeline}} when
+\code{pipeline} is NULL.}
+
+\item{width,height}{Integers, divisible by 16.}
+
+\item{prompt_embeds,pooled_prompt_embeds}{Optional precomputed text
+embeddings (skip the text encoders).}
+}
+\value{
+Invisibly, \code{list(image, metadata)} where \code{image} is
+ an [H, W, 3] array in [0, 1].
+}
+\description{
+4-step distilled text-to-image generation (no classifier-free
+guidance): T5 + CLIP prompt encoding, flow-matching Euler denoising
+over the packed latent sequence, and 16-channel VAE decode. With
+phase offloading each component is the sole GPU tenant for its phase.
+}
diff --git a/man/unigram_tokenizer.Rd b/man/unigram_tokenizer.Rd
new file mode 100644
index 0000000..4adf0f4
--- /dev/null
+++ b/man/unigram_tokenizer.Rd
@@ -0,0 +1,17 @@
+% tinyrox says don't edit this manually, but it can't stop you!
+\name{unigram_tokenizer}
+\alias{unigram_tokenizer}
+\title{Load a Unigram tokenizer from tokenizer.json}
+\usage{
+unigram_tokenizer(tokenizer_path)
+}
+\arguments{
+\item{tokenizer_path}{Path to a HuggingFace tokenizer.json with a
+Unigram model, or a directory containing one.}
+}
+\value{
+A \code{unigram_tokenizer} object.
+}
+\description{
+Load a Unigram tokenizer from tokenizer.json
+}
diff --git a/man/vae_decoder_native.Rd b/man/vae_decoder_native.Rd
index 5dfc6aa..c6823f8 100644
--- a/man/vae_decoder_native.Rd
+++ b/man/vae_decoder_native.Rd
@@ -3,12 +3,20 @@
\alias{vae_decoder_native}
\title{Native VAE Decoder}
\usage{
-vae_decoder_native(latent_channels = 4, out_channels = 3)
+vae_decoder_native(latent_channels = 4, out_channels = 3,
+ block_channels = c(512, 512, 256, 128), norm_groups = 32)
}
\arguments{
-\item{latent_channels}{Number of latent channels (default 4)}
+\item{latent_channels}{Number of latent channels (4 for SD/SDXL,
+16 for FLUX/SD3)}
\item{out_channels}{Number of output channels (default 3 for RGB)}
+
+\item{block_channels}{Decoder block channels (reversed encoder
+block_out_channels; default matches SD/SDXL and FLUX)}
+
+\item{norm_groups}{Group norm groups (default 32; must divide every
+entry of \code{block_channels})}
}
\value{
An nn_module representing the VAE decoder
diff --git a/tools/compare_translation.R b/tools/compare_translation.R
new file mode 100644
index 0000000..87feaad
--- /dev/null
+++ b/tools/compare_translation.R
@@ -0,0 +1,192 @@
+# Structural translation audit for the FLUX port, using treesitR.
+#
+# Parses the Python reference (diffusers / transformers, Apache-2.0) and
+# the R port, extracts numeric literals and torch-op call names per
+# pairing, and reports the set differences. The point is a short,
+# reviewable drift report - not a zero-diff goal: R-side plumbing
+# (seq_len, message, ...) and Python-side plumbing (isinstance, super,
+# ...) are filtered, the rest is eyeballed.
+#
+# Run from the package root: r tools/compare_translation.R
+# Requires: treesitR, the ref/ symlink, and tools/cache/modeling_t5.py
+# (curl it from transformers v4.49.0 if absent).
+
+library(treesitR)
+
+DIFFUSERS <- "ref/upstream/diffusers/src/diffusers"
+
+# ---- tree helpers -------------------------------------------------------------
+
+parse_file <- function(path, language) {
+ p <- ts_parser_new()
+ ts_parser_set_language(p, language)
+ src <- paste(readLines(path, warn = FALSE), collapse = "\n")
+ ts_tree_root_node(ts_parse(p, src))
+}
+
+walk_collect <- function(node, fn) {
+ acc <- list()
+ recurse <- function(n) {
+ r <- fn(n)
+ if (!is.null(r)) {
+ acc[[length(acc) + 1L]] <<- r
+ }
+ for (child in ts_node_children(n, named = TRUE)) {
+ recurse(child)
+ }
+ }
+ recurse(node)
+ acc
+}
+
+# Named scopes (class_definition / function_definition) from a Python tree
+py_scopes <- function(root, names) {
+ walk_collect(root, function(n) {
+ if (ts_node_type(n) %in% c("class_definition", "function_definition")) {
+ kids <- ts_node_children(n, named = TRUE)
+ if (length(kids) && ts_node_type(kids[[1]]) == "identifier" &&
+ ts_node_text(kids[[1]]) %in% names) {
+ return(n)
+ }
+ }
+ NULL
+ })
+}
+
+# Numeric literal values under a node (sign-insensitive)
+literals <- function(node) {
+ vals <- walk_collect(node, function(n) {
+ if (ts_node_type(n) %in% c("integer", "float")) {
+ txt <- gsub("L$|_", "", ts_node_text(n))
+ v <- suppressWarnings(as.numeric(txt))
+ if (!is.na(v)) {
+ return(v)
+ }
+ }
+ NULL
+ })
+ sort(unique(unlist(vals)))
+}
+
+# Canonical call names under a node: last identifier of the callee,
+# stripped of torch_/nnf_ prefixes
+call_names <- function(node) {
+ names <- walk_collect(node, function(n) {
+ if (ts_node_type(n) != "call") {
+ return(NULL)
+ }
+ callee <- ts_node_children(n, named = TRUE)[[1]]
+ if (ts_node_type(callee) %in%
+ c("attribute", "extract_operator", "namespace_operator")) {
+ kids <- ts_node_children(callee, named = TRUE)
+ callee <- kids[[length(kids)]]
+ }
+ if (ts_node_type(callee) != "identifier") {
+ return(NULL)
+ }
+ sub("^torch_|^nnf_", "", ts_node_text(callee))
+ })
+ sort(unique(unlist(names)))
+}
+
+# Plumbing that legitimately exists on only one side
+PY_NOISE <- c(
+ "super", "range", "len", "isinstance", "hasattr", "getattr", "print",
+ "int", "float", "str", "list", "dict", "tuple", "set", "zip",
+ "enumerate", "ValueError", "ImportError", "warning", "warn",
+ "deprecate", "items", "keys", "values", "pop", "update", "get",
+ "append", "join", "startswith", "endswith", "register_to_config",
+ "apply_lora_scale", "maybe_allow_in_graph", "ceil", "is_grad_enabled",
+ "_gradient_checkpointing_func", "register_buffer", "ModuleList",
+ "Module", "signature", "is_torch_npu_available",
+ "maybe_adjust_dtype_for_device", "dispatch_attention_fn",
+ "set_processor", "processor", "retrieve_timesteps", "register_modules",
+ "randn_tensor", "postprocess", "numpy", "maybe_free_model_hooks",
+ "MultiPipelineCallbacks", "PipelineCallback", "progress_bar"
+)
+R_NOISE <- c(
+ "c", "list", "seq_len", "seq_along", "lapply", "vapply", "function",
+ "paste0", "paste", "message", "stop", "warning", "sprintf", "length",
+ "names", "file.path", "is.null", "invisible", "structure", "inherits",
+ "getOption", "options", "requireNamespace", "Sys.time", "difftime",
+ "Sys.getenv", "Sys.setenv", "close", "txtProgressBar",
+ "setTxtProgressBar", "nn_module", "nn_module_list", "gc", "rm",
+ "isTRUE", "as.integer", "as.numeric", "nzchar", "nchar", "dirname",
+ "path.expand", "file.exists", "dir.exists", "fromJSON", "tryCatch",
+ "match.arg", "strrep", "modifyList", "do.call", "Filter", "Negate",
+ "grepl", "sub", "gsub", "startsWith", "endsWith", "setdiff", "head",
+ "utils", "hub_download", "filename_from_prompt", "save_image",
+ "clear_vram", "onload", "offload", "print"
+)
+
+report_pair <- function(label, py_nodes, r_files) {
+ r_lits <- c()
+ r_calls <- c()
+ for (f in r_files) {
+ root <- parse_file(f, ts_language_r())
+ r_lits <- union(r_lits, literals(root))
+ r_calls <- union(r_calls, call_names(root))
+ }
+ py_lits <- sort(unique(unlist(lapply(py_nodes, literals))))
+ py_calls <- sort(unique(unlist(lapply(py_nodes, call_names))))
+
+ cat("\n== ", label, " ==\n", sep = "")
+ cat("literals only in Python: ",
+ paste(setdiff(py_lits, r_lits), collapse = " "), "\n")
+ cat("literals only in R: ",
+ paste(setdiff(r_lits, py_lits), collapse = " "), "\n")
+ cat("calls only in Python: ",
+ paste(setdiff(setdiff(py_calls, r_calls), PY_NOISE), collapse = " "),
+ "\n")
+ cat("calls only in R: ",
+ paste(setdiff(setdiff(r_calls, py_calls), R_NOISE), collapse = " "),
+ "\n")
+}
+
+# ---- pairings -------------------------------------------------------------------
+
+# 1. Transformer stack: blocks, norms, embedders, RoPE
+tf <- parse_file(file.path(DIFFUSERS, "models/transformers/transformer_flux.py"),
+ ts_language_python())
+norm <- parse_file(file.path(DIFFUSERS, "models/normalization.py"),
+ ts_language_python())
+emb <- parse_file(file.path(DIFFUSERS, "models/embeddings.py"),
+ ts_language_python())
+act <- parse_file(file.path(DIFFUSERS, "models/attention.py"),
+ ts_language_python())
+py_transformer <- c(
+ py_scopes(tf, c("FluxAttnProcessor", "FluxAttention",
+ "FluxSingleTransformerBlock", "FluxTransformerBlock",
+ "FluxPosEmbed", "FluxTransformer2DModel",
+ "_get_qkv_projections")),
+ py_scopes(norm, c("AdaLayerNormZero", "AdaLayerNormZeroSingle",
+ "AdaLayerNormContinuous")),
+ py_scopes(emb, c("get_1d_rotary_pos_embed", "apply_rotary_emb",
+ "CombinedTimestepTextProjEmbeddings",
+ "PixArtAlphaTextProjection", "Timesteps",
+ "TimestepEmbedding", "get_timestep_embedding")),
+ py_scopes(act, c("FeedForward"))
+)
+report_pair("transformer stack",
+ py_transformer,
+ c("R/dit_flux_modules.R", "R/dit_flux.R", "R/rope_flux.R"))
+
+# 2. Pipeline flow
+pipe <- parse_file(file.path(DIFFUSERS, "pipelines/flux/pipeline_flux.py"),
+ ts_language_python())
+report_pair("pipeline",
+ py_scopes(pipe, c("FluxPipeline", "calculate_shift")),
+ c("R/txt2img_flux.R"))
+
+# 3. T5 encoder
+t5_path <- "tools/cache/modeling_t5.py"
+if (file.exists(t5_path)) {
+ t5 <- parse_file(t5_path, ts_language_python())
+ report_pair("t5 encoder",
+ py_scopes(t5, c("T5LayerNorm", "T5DenseGatedActDense",
+ "T5Attention", "T5LayerSelfAttention", "T5LayerFF",
+ "T5Block", "T5Stack", "T5EncoderModel")),
+ c("R/t5_text_encoder.R"))
+} else {
+ cat("\n(t5 encoder pairing skipped: tools/cache/modeling_t5.py missing)\n")
+}
diff --git a/tools/gen_fixtures.sh b/tools/gen_fixtures.sh
index e3f2f73..83b66a6 100755
--- a/tools/gen_fixtures.sh
+++ b/tools/gen_fixtures.sh
@@ -19,6 +19,7 @@ for s in "${scripts[@]}"; do
--with torch \
--with numpy \
--with safetensors \
+ --with transformers \
--with huggingface_hub \
--with packaging \
--with filelock \
diff --git a/tools/gen_fixtures_flux.py b/tools/gen_fixtures_flux.py
new file mode 100644
index 0000000..c72e788
--- /dev/null
+++ b/tools/gen_fixtures_flux.py
@@ -0,0 +1,70 @@
+# Generate FLUX parity fixtures for the R port.
+#
+# Runs the diffusers reference implementation (Apache-2.0) on small fixed
+# inputs and saves {inputs, expected outputs} as safetensors fixtures that
+# the R tinytest suite compares against. Run via tools/gen_fixtures.sh;
+# never executed at package test/run time.
+
+import os
+import sys
+
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src"))
+
+from diffusers.models.embeddings import apply_rotary_emb # noqa: E402
+from diffusers.models.transformers.transformer_flux import FluxPosEmbed # noqa: E402
+from diffusers.pipelines.flux.pipeline_flux import FluxPipeline # noqa: E402
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures")
+os.makedirs(OUT_DIR, exist_ok=True)
+
+torch.manual_seed(42)
+fx = {}
+
+# --- latent image ids ---------------------------------------------------------
+# Asymmetric grid (8 rows x 12 cols) so a row/col swap fails the test.
+img_ids = FluxPipeline._prepare_latent_image_ids(
+ batch_size=1, height=8, width=12, device="cpu", dtype=torch.float32
+)
+fx["img_ids"] = img_ids # [96, 3]
+
+# --- FluxPosEmbed: full-size axes (16, 56, 56) --------------------------------
+txt_ids = torch.zeros(7, 3)
+ids = torch.cat([txt_ids, img_ids], dim=0) # [103, 3]
+fx["ids"] = ids
+
+pos_full = FluxPosEmbed(theta=10000, axes_dim=[16, 56, 56])
+cos_full, sin_full = pos_full(ids)
+fx["pos_full_cos"] = cos_full # [103, 128]
+fx["pos_full_sin"] = sin_full
+
+# --- FluxPosEmbed: tiny axes (2, 2, 4) used by the block-level fixtures -------
+pos_tiny = FluxPosEmbed(theta=10000, axes_dim=[2, 2, 4])
+cos_tiny, sin_tiny = pos_tiny(ids)
+fx["pos_tiny_cos"] = cos_tiny # [103, 8]
+fx["pos_tiny_sin"] = sin_tiny
+
+# --- apply_rotary_emb on [B, H, S, D] (sequence_dim=2, unbind_dim=-1) ----------
+B, H, S, D = 2, 4, 103, 8
+x = torch.randn(B, H, S, D)
+fx["rot_x"] = x
+fx["rot_out"] = apply_rotary_emb(x, (cos_tiny, sin_tiny), sequence_dim=2)
+
+# fp16 dtype preservation through apply
+x16 = torch.randn(B, H, S, D, dtype=torch.float16)
+fx["rot_x_f16"] = x16
+fx["rot_out_f16"] = apply_rotary_emb(x16, (cos_tiny, sin_tiny), sequence_dim=2)
+
+# --- latent pack / unpack -------------------------------------------------------
+# Latent [2, 16, 8, 12] corresponds to a 64x96 pixel image (vae_scale_factor 8).
+lat = torch.randn(2, 16, 8, 12)
+packed = FluxPipeline._pack_latents(lat, 2, 16, 8, 12)
+fx["pack_in"] = lat
+fx["pack_out"] = packed # [2, 24, 64]
+fx["unpack_out"] = FluxPipeline._unpack_latents(packed, height=64, width=96, vae_scale_factor=8)
+
+fx = {k: v.contiguous() for k, v in fx.items()}
+save_file(fx, os.path.join(OUT_DIR, "rope_flux.safetensors"))
+print(f"wrote {len(fx)} tensors to {OUT_DIR}/rope_flux.safetensors")
diff --git a/tools/gen_fixtures_flux_clip_vae.py b/tools/gen_fixtures_flux_clip_vae.py
new file mode 100644
index 0000000..cb5e80e
--- /dev/null
+++ b/tools/gen_fixtures_flux_clip_vae.py
@@ -0,0 +1,94 @@
+# Generate CLIP (quick_gelu + pooled output) and 16-channel VAE decoder
+# parity fixtures for the FLUX R port.
+#
+# Uses HF transformers CLIPTextModel and diffusers AutoencoderKL
+# (Apache-2.0) with tiny random-init configs. Run via
+# tools/gen_fixtures.sh; never executed at package test/run time.
+
+import os
+import sys
+
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src"))
+
+from transformers import CLIPTextConfig, CLIPTextModel # noqa: E402
+
+from diffusers import AutoencoderKL # noqa: E402
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures")
+os.makedirs(OUT_DIR, exist_ok=True)
+
+torch.manual_seed(46)
+
+# --- tiny CLIP text model (quick_gelu, legacy argmax pooling) ------------------
+# eos_token_id=2 selects the legacy argmax(input_ids) pooling path, which
+# is what FLUX's CLIP-L config uses.
+clip_cfg = CLIPTextConfig(
+ vocab_size=1000,
+ hidden_size=16,
+ intermediate_size=32,
+ num_hidden_layers=2,
+ num_attention_heads=2,
+ max_position_embeddings=77,
+ hidden_act="quick_gelu",
+ eos_token_id=2,
+)
+clip = CLIPTextModel(clip_cfg)
+clip.eval()
+
+# EOS (as the max id, 999) at different positions; padding after
+input_ids = torch.tensor([
+ [49, 23, 61, 7, 999, 0, 0, 0],
+ [33, 999, 5, 5, 5, 5, 5, 5],
+])
+with torch.no_grad():
+ out = clip(input_ids, output_hidden_states=False)
+
+save_file(
+ {k: v.contiguous() for k, v in clip.state_dict().items()},
+ os.path.join(OUT_DIR, "clip_tiny.safetensors"),
+)
+
+io = {
+ "clip_input_ids": input_ids,
+ "clip_last_hidden": out.last_hidden_state,
+ "clip_pooled": out.pooler_output,
+}
+
+# --- tiny 16-channel VAE (FLUX/SD3 shape: no quant convs, mid attention) --------
+vae = AutoencoderKL(
+ in_channels=3,
+ out_channels=3,
+ down_block_types=("DownEncoderBlock2D",) * 4,
+ up_block_types=("UpDecoderBlock2D",) * 4,
+ block_out_channels=(8, 16, 32, 32),
+ layers_per_block=2,
+ latent_channels=16,
+ norm_num_groups=8,
+ use_quant_conv=False,
+ use_post_quant_conv=False,
+ mid_block_add_attention=True,
+ sample_size=32,
+)
+vae.eval()
+
+latent = torch.randn(1, 16, 4, 4)
+with torch.no_grad():
+ image = vae.decode(latent, return_dict=False)[0]
+
+save_file(
+ {k: v.contiguous() for k, v in vae.state_dict().items()},
+ os.path.join(OUT_DIR, "vae16_tiny.safetensors"),
+)
+
+io["vae_latent"] = latent
+io["vae_image"] = image
+io = {k: v.contiguous() for k, v in io.items()}
+# Metadata shifts the header size: without it this file's leading bytes
+# happen to match `file`'s VAX COFF magic and R CMD check flags it as
+# an executable
+save_file(io, os.path.join(OUT_DIR, "clip_vae_io.safetensors"),
+ metadata={"purpose": "diffuseR FLUX test fixture"})
+print(f"wrote clip_tiny, vae16_tiny, and {len(io)} io tensors to {OUT_DIR}")
diff --git a/tools/gen_fixtures_flux_dit.py b/tools/gen_fixtures_flux_dit.py
new file mode 100644
index 0000000..aace645
--- /dev/null
+++ b/tools/gen_fixtures_flux_dit.py
@@ -0,0 +1,154 @@
+# Generate FLUX transformer-block parity fixtures for the R port.
+#
+# Runs the diffusers reference modules (Apache-2.0) with tiny random-init
+# configs and saves {state dicts, inputs, outputs} as safetensors fixtures
+# for the R tinytest suite. Run via tools/gen_fixtures.sh; never executed
+# at package test/run time.
+
+import os
+import sys
+
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src"))
+
+from diffusers.models.normalization import ( # noqa: E402
+ AdaLayerNormContinuous,
+ AdaLayerNormZero,
+ AdaLayerNormZeroSingle,
+)
+from diffusers.models.transformers.transformer_flux import ( # noqa: E402
+ FluxAttention,
+ FluxPosEmbed,
+ FluxSingleTransformerBlock,
+ FluxTransformerBlock,
+)
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures")
+os.makedirs(OUT_DIR, exist_ok=True)
+
+torch.manual_seed(43)
+fx = {}
+
+# Tiny config: dim 16 = 2 heads x head_dim 8, rope axes (2, 2, 4).
+DIM, HEADS, HEAD_DIM = 16, 2, 8
+B, S_TXT = 2, 7
+GRID_H, GRID_W = 8, 12
+S_IMG = GRID_H * GRID_W # 96
+S = S_TXT + S_IMG # 103
+
+# Position ids and rotary freqs shared by all attention fixtures
+img_ids = torch.zeros(GRID_H, GRID_W, 3)
+img_ids[..., 1] = torch.arange(GRID_H)[:, None]
+img_ids[..., 2] = torch.arange(GRID_W)[None, :]
+img_ids = img_ids.reshape(S_IMG, 3)
+ids = torch.cat([torch.zeros(S_TXT, 3), img_ids], dim=0)
+rope_cos, rope_sin = FluxPosEmbed(theta=10000, axes_dim=[2, 2, 4])(ids)
+fx["rope_cos"] = rope_cos
+fx["rope_sin"] = rope_sin
+
+img_x = torch.randn(B, S_IMG, DIM)
+txt_x = torch.randn(B, S_TXT, DIM)
+joint_x = torch.randn(B, S, DIM)
+temb = torch.randn(B, DIM)
+fx["img_x"] = img_x
+fx["txt_x"] = txt_x
+fx["joint_x"] = joint_x
+fx["temb"] = temb
+
+
+def add_state(prefix, module):
+ for k, v in module.state_dict().items():
+ fx[f"{prefix}.{k}"] = v
+
+
+# --- AdaLayerNormZero -----------------------------------------------------------
+adazero = AdaLayerNormZero(DIM)
+add_state("adazero", adazero)
+with torch.no_grad():
+ x_norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = adazero(joint_x, emb=temb)
+fx["adazero_x_norm"] = x_norm
+fx["adazero_gate_msa"] = gate_msa
+fx["adazero_shift_mlp"] = shift_mlp
+fx["adazero_scale_mlp"] = scale_mlp
+fx["adazero_gate_mlp"] = gate_mlp
+
+# --- AdaLayerNormZeroSingle -------------------------------------------------------
+adasingle = AdaLayerNormZeroSingle(DIM)
+add_state("adasingle", adasingle)
+with torch.no_grad():
+ xs_norm, gate = adasingle(joint_x, emb=temb)
+fx["adasingle_x_norm"] = xs_norm
+fx["adasingle_gate"] = gate
+
+# --- AdaLayerNormContinuous (norm_out config) --------------------------------------
+adacont = AdaLayerNormContinuous(DIM, DIM, elementwise_affine=False, eps=1e-6)
+add_state("adacont", adacont)
+with torch.no_grad():
+ fx["adacont_out"] = adacont(joint_x, temb)
+
+# --- FluxAttention, double-stream variant (added_kv, joint attention) ---------------
+attn_d = FluxAttention(
+ query_dim=DIM,
+ added_kv_proj_dim=DIM,
+ dim_head=HEAD_DIM,
+ heads=HEADS,
+ out_dim=DIM,
+ context_pre_only=False,
+ bias=True,
+ eps=1e-6,
+)
+add_state("attnd", attn_d)
+with torch.no_grad():
+ attn_out, ctx_out = attn_d(
+ hidden_states=img_x,
+ encoder_hidden_states=txt_x,
+ image_rotary_emb=(rope_cos, rope_sin),
+ )
+fx["attnd_out"] = attn_out
+fx["attnd_ctx_out"] = ctx_out
+
+# --- FluxAttention, pre_only variant (single blocks) ---------------------------------
+attn_s = FluxAttention(
+ query_dim=DIM,
+ dim_head=HEAD_DIM,
+ heads=HEADS,
+ out_dim=DIM,
+ bias=True,
+ eps=1e-6,
+ pre_only=True,
+)
+add_state("attns", attn_s)
+with torch.no_grad():
+ fx["attns_out"] = attn_s(hidden_states=joint_x, image_rotary_emb=(rope_cos, rope_sin))
+
+# --- FluxTransformerBlock (double) ----------------------------------------------------
+dbl = FluxTransformerBlock(dim=DIM, num_attention_heads=HEADS, attention_head_dim=HEAD_DIM)
+add_state("dbl", dbl)
+with torch.no_grad():
+ enc_out, hid_out = dbl(
+ hidden_states=img_x,
+ encoder_hidden_states=txt_x,
+ temb=temb,
+ image_rotary_emb=(rope_cos, rope_sin),
+ )
+fx["dbl_enc_out"] = enc_out
+fx["dbl_hid_out"] = hid_out
+
+# --- FluxSingleTransformerBlock --------------------------------------------------------
+sgl = FluxSingleTransformerBlock(dim=DIM, num_attention_heads=HEADS, attention_head_dim=HEAD_DIM)
+add_state("sgl", sgl)
+with torch.no_grad():
+ s_enc_out, s_hid_out = sgl(
+ hidden_states=img_x,
+ encoder_hidden_states=txt_x,
+ temb=temb,
+ image_rotary_emb=(rope_cos, rope_sin),
+ )
+fx["sgl_enc_out"] = s_enc_out
+fx["sgl_hid_out"] = s_hid_out
+
+fx = {k: v.contiguous() for k, v in fx.items()}
+save_file(fx, os.path.join(OUT_DIR, "dit_flux.safetensors"))
+print(f"wrote {len(fx)} tensors to {OUT_DIR}/dit_flux.safetensors")
diff --git a/tools/gen_fixtures_flux_model.py b/tools/gen_fixtures_flux_model.py
new file mode 100644
index 0000000..634b985
--- /dev/null
+++ b/tools/gen_fixtures_flux_model.py
@@ -0,0 +1,84 @@
+# Generate a tiny random-init FluxTransformer2DModel parity fixture for
+# the R port, plus a sharded save_pretrained checkpoint used by the
+# checkpoint-loading and quantization tests.
+#
+# Uses the diffusers reference implementation (Apache-2.0). Run via
+# tools/gen_fixtures.sh; never executed at package test/run time.
+
+import os
+import shutil
+import sys
+
+import torch
+from safetensors.torch import save_file
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "ref", "upstream", "diffusers", "src"))
+
+from diffusers import FluxTransformer2DModel # noqa: E402
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures")
+CKPT_DIR = os.path.join(OUT_DIR, "flux_tiny_ckpt")
+os.makedirs(OUT_DIR, exist_ok=True)
+
+torch.manual_seed(44)
+
+model = FluxTransformer2DModel(
+ patch_size=1,
+ in_channels=4,
+ num_layers=1,
+ num_single_layers=1,
+ attention_head_dim=8,
+ num_attention_heads=2,
+ joint_attention_dim=16,
+ pooled_projection_dim=12,
+ axes_dims_rope=(2, 2, 4),
+)
+model.eval()
+
+B, S_TXT = 2, 7
+GRID_H, GRID_W = 8, 12
+S_IMG = GRID_H * GRID_W
+
+hidden = torch.randn(B, S_IMG, 4)
+encoder = torch.randn(B, S_TXT, 16)
+pooled = torch.randn(B, 12)
+timestep = torch.tensor([1.0, 0.25]) # sigma space; model multiplies by 1000
+txt_ids = torch.zeros(S_TXT, 3)
+img_ids = torch.zeros(GRID_H, GRID_W, 3)
+img_ids[..., 1] = torch.arange(GRID_H)[:, None]
+img_ids[..., 2] = torch.arange(GRID_W)[None, :]
+img_ids = img_ids.reshape(S_IMG, 3)
+
+with torch.no_grad():
+ out = model(
+ hidden_states=hidden,
+ encoder_hidden_states=encoder,
+ pooled_projections=pooled,
+ timestep=timestep,
+ txt_ids=txt_ids,
+ img_ids=img_ids,
+ return_dict=False,
+ )[0]
+
+fx = {f"model.{k}": v for k, v in model.state_dict().items()}
+fx.update(
+ {
+ "hidden": hidden,
+ "encoder": encoder,
+ "pooled": pooled,
+ "timestep": timestep,
+ "txt_ids": txt_ids,
+ "img_ids": img_ids,
+ "out": out,
+ }
+)
+fx = {k: v.contiguous() for k, v in fx.items()}
+save_file(fx, os.path.join(OUT_DIR, "flux_model.safetensors"))
+print(f"wrote {len(fx)} tensors to {OUT_DIR}/flux_model.safetensors")
+
+# Sharded checkpoint in the real HF layout (config.json + shards +
+# diffusion_pytorch_model.safetensors.index.json) for loader tests
+if os.path.isdir(CKPT_DIR):
+ shutil.rmtree(CKPT_DIR)
+model.save_pretrained(CKPT_DIR, max_shard_size="30KB")
+print(f"wrote sharded checkpoint to {CKPT_DIR}: {sorted(os.listdir(CKPT_DIR))}")
diff --git a/tools/gen_fixtures_t5.py b/tools/gen_fixtures_t5.py
new file mode 100644
index 0000000..49b143c
--- /dev/null
+++ b/tools/gen_fixtures_t5.py
@@ -0,0 +1,66 @@
+# Generate T5 encoder parity fixtures for the R port.
+#
+# Uses HF transformers' T5EncoderModel (Apache-2.0) with a tiny
+# random-init config, plus an exact integer fixture for the relative
+# position bucketing. Run via tools/gen_fixtures.sh; never executed at
+# package test/run time.
+
+import os
+import shutil
+
+import torch
+from safetensors.torch import save_file
+from transformers import T5Config, T5EncoderModel
+from transformers.models.t5.modeling_t5 import T5Attention
+
+OUT_DIR = os.path.join(os.path.dirname(__file__), "..", "inst", "tinytest", "fixtures")
+CKPT_DIR = os.path.join(OUT_DIR, "t5_tiny_ckpt")
+os.makedirs(OUT_DIR, exist_ok=True)
+
+torch.manual_seed(45)
+fx = {}
+
+# --- relative position bucketing (exact integer parity) -----------------------
+rel = torch.arange(-200, 201).unsqueeze(0)
+buckets = T5Attention._relative_position_bucket(
+ rel, bidirectional=True, num_buckets=32, max_distance=128
+)
+fx["rel_positions"] = rel
+fx["rel_buckets"] = buckets
+
+# --- tiny encoder ---------------------------------------------------------------
+config = T5Config(
+ vocab_size=100,
+ d_model=16,
+ d_kv=4,
+ num_heads=4,
+ d_ff=32,
+ num_layers=2,
+ feed_forward_proj="gated-gelu",
+ relative_attention_num_buckets=32,
+ relative_attention_max_distance=128,
+ layer_norm_epsilon=1e-6,
+ dropout_rate=0.0,
+)
+model = T5EncoderModel(config)
+model.eval()
+
+# Padded batch, and (matching FLUX) NO attention mask
+input_ids = torch.tensor([
+ [5, 23, 61, 7, 19, 88, 42, 3, 1, 0, 0, 0],
+ [33, 14, 2, 71, 1, 0, 0, 0, 0, 0, 0, 0],
+])
+with torch.no_grad():
+ out = model(input_ids=input_ids).last_hidden_state
+
+fx["input_ids"] = input_ids
+fx["out"] = out
+fx = {k: v.contiguous() for k, v in fx.items()}
+save_file(fx, os.path.join(OUT_DIR, "t5_flux.safetensors"))
+print(f"wrote {len(fx)} tensors to {OUT_DIR}/t5_flux.safetensors")
+
+# Checkpoint dir in the transformers layout for the loader test
+if os.path.isdir(CKPT_DIR):
+ shutil.rmtree(CKPT_DIR)
+model.save_pretrained(CKPT_DIR)
+print(f"wrote checkpoint to {CKPT_DIR}: {sorted(os.listdir(CKPT_DIR))}")
diff --git a/tools/gen_t5_tokenizer_cases.py b/tools/gen_t5_tokenizer_cases.py
new file mode 100644
index 0000000..4db303b
--- /dev/null
+++ b/tools/gen_t5_tokenizer_cases.py
@@ -0,0 +1,117 @@
+# Generate T5 Unigram tokenizer parity cases for the R port.
+#
+# Uses FLUX.1-schnell's shipped tokenizer_2/tokenizer.json from the
+# HuggingFace cache when available (the authoritative artifact:
+# Metaspace prepend_scheme "always"); otherwise converts spiece.model
+# from the public google/t5-v1_1-xxl repo and patches the prepend scheme
+# to match. Writes:
+# - tools/cache/tokenizer_t5.json (dev copy, gitignored)
+# - inst/tinytest/fixtures/t5_tokenizer_cases.json (checked in)
+#
+# Run:
+# uv run --no-project --with transformers --with sentencepiece \
+# --with protobuf --with torch --index https://download.pytorch.org/whl/cpu \
+# --index-strategy unsafe-best-match python tools/gen_t5_tokenizer_cases.py
+
+import json
+import os
+
+from huggingface_hub import hf_hub_download
+from transformers import PreTrainedTokenizerFast
+
+ROOT = os.path.join(os.path.dirname(__file__), "..")
+CACHE_DIR = os.path.join(ROOT, "tools", "cache")
+FIXTURE = os.path.join(ROOT, "inst", "tinytest", "fixtures", "t5_tokenizer_cases.json")
+os.makedirs(CACHE_DIR, exist_ok=True)
+TOK_JSON = os.path.join(CACHE_DIR, "tokenizer_t5.json")
+
+def is_real_tokenizer(path):
+ try:
+ with open(path) as f:
+ return len(json.load(f)["model"]["vocab"]) > 30000
+ except Exception:
+ return False
+
+
+if not is_real_tokenizer(TOK_JSON):
+ import glob
+
+ shipped = glob.glob(os.path.expanduser(
+ "~/.cache/huggingface/hub/models--black-forest-labs--FLUX.1-schnell/"
+ "snapshots/*/tokenizer_2/tokenizer.json"
+ ))
+ if shipped:
+ with open(shipped[0], "rb") as f_in, open(TOK_JSON, "wb") as f_out:
+ f_out.write(f_in.read())
+ print(f"using shipped tokenizer: {shipped[0]}")
+ else:
+ # Conversion fallback: requires transformers<5 (v5 dropped the
+ # sentencepiece slow->fast converter and emits a 104-token stub)
+ from transformers import T5TokenizerFast
+
+ spiece = hf_hub_download("google/t5-v1_1-xxl", "spiece.model")
+ T5TokenizerFast(vocab_file=spiece).backend_tokenizer.save(TOK_JSON)
+ tj = json.load(open(TOK_JSON))
+ assert len(tj["model"]["vocab"]) > 30000, \
+ "conversion produced a stub vocab; pin transformers<5"
+ for part in ("pre_tokenizer", "decoder"):
+ if tj.get(part, {}).get("type") == "Metaspace":
+ tj[part]["prepend_scheme"] = "always"
+ json.dump(tj, open(TOK_JSON, "w"), ensure_ascii=False)
+ print("using converted spiece.model, prepend_scheme patched to 'always'")
+else:
+ print(f"using existing {TOK_JSON}")
+
+tok = PreTrainedTokenizerFast(
+ tokenizer_file=TOK_JSON, pad_token="", eos_token="",
+ unk_token="",
+)
+
+PROMPTS = [
+ "a photo of a cat",
+ "A sunset over mountains, ultra detailed, 8k",
+ "Hello, world!",
+ "The quick brown fox jumps over the lazy dog.",
+ "it's a beautiful day; isn't it?",
+ "3.14159 and 2,000,000 dollars",
+ "state-of-the-art text-to-image generation",
+ " leading spaces",
+ "trailing spaces ",
+ "double and triple spaces",
+ "UPPERCASE lowercase MiXeD",
+ "email@example.com and https://example.org/path?q=1",
+ 'quotes "double" and \'single\'',
+ "(parentheses) [brackets] {braces}",
+ "semi;colon co:lon sla/sh back\\slash",
+ "underscores_and_snake_case",
+ "a",
+ "",
+ "café résumé naïve",
+ "em—dash and en–dash",
+ "100% of $50 + €20",
+ "An astronaut riding a horse on Mars, photorealistic",
+ "watercolor painting of a fox in a snowy forest",
+ "the mitochondria is the powerhouse of the cell",
+ ("The transformer architecture uses self-attention mechanisms "
+ "to model long-range dependencies in sequences. ") * 12,
+]
+
+cases = []
+for text in PROMPTS:
+ ids = tok(text, add_special_tokens=True)["input_ids"]
+ cases.append({"text": text, "ids": ids})
+
+# Padding/truncation behavior at a small max_length
+padded = []
+for text in [PROMPTS[0], PROMPTS[3], PROMPTS[24]]:
+ enc = tok(text, max_length=16, padding="max_length", truncation=True)
+ padded.append({
+ "text": text,
+ "max_length": 16,
+ "ids": enc["input_ids"],
+ "mask": enc["attention_mask"],
+ })
+
+with open(FIXTURE, "w", encoding="utf-8") as f:
+ json.dump({"cases": cases, "padded": padded}, f, ensure_ascii=False, indent=1)
+print(f"wrote {len(cases)} cases + {len(padded)} padded cases to {FIXTURE}")